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

//! Types supporting documentation for code written in or for Starlark.

// TODO(nga): document it
#![allow(missing_docs)]

mod markdown;

use std::collections::HashMap;

use allocative::Allocative;
use dupe::Dupe;
use itertools::Itertools;
pub use markdown::MarkdownFlavor;
pub use markdown::RenderMarkdown;
use once_cell::sync::Lazy;
use regex::Regex;
use regex::RegexBuilder;
use serde::Serialize;
pub use starlark_derive::StarlarkDocs;
use starlark_map::small_map::SmallMap;

use crate as starlark;
use crate::codemap::Spanned;
use crate::syntax::ast::AstLiteral;
use crate::syntax::ast::AstPayload;
use crate::syntax::ast::AstStmtP;
use crate::syntax::ast::ExprP;
use crate::syntax::ast::StmtP;
use crate::typing::Ty;
use crate::values::StarlarkValue;
use crate::values::Trace;
use crate::values::Value;

/// There have been bugs around line endings in the textwrap crate. Just join
/// into a single string, and trim the line endings.
fn wrap_trimmed(s: &str, width: usize) -> String {
    textwrap::wrap(s, width).join("\n").trim_end().to_owned()
}

/// There have been bugs around line endings in the textwrap crate. Just trim the line endings.
fn indent_trimmed(s: &str, prefix: &str) -> String {
    textwrap::indent(s, prefix).trim_end().to_owned()
}

/// The documentation provided by a user for a specific module, object, function, etc.
#[derive(Debug, Clone, PartialEq, Serialize, Trace, Default, Allocative)]
pub struct DocString {
    /// The first line of a doc string. This has whitespace trimmed from it.
    pub summary: String,
    /// The contents of a doc string that follow the summary, and a single blank line.
    /// This also has whitespace trimmed from it, and it is dedented.
    pub details: Option<String>,
}

impl DocString {
    /// Render this docstring as a "starlark" docstring.
    fn render_as_code(&self) -> String {
        let s = match &self.details {
            Some(details) => {
                format!("{}\n\n{}", self.summary, details)
            }
            None => self.summary.clone(),
        };
        wrap_trimmed(&s, 80)
    }

    /// Render the docstring as in `render_as_code`, but surround it in triple quotes,
    /// a common convention in starlark docstrings.
    fn render_as_quoted_code(&self) -> String {
        format!("\"\"\"\n{}\n\"\"\"", self.render_as_code())
    }
}

/// Controls the formatting to use when parsing `DocString`s from raw docstrings
#[derive(Copy, Clone, Dupe)]
pub enum DocStringKind {
    /// Docstrings provided by users in starlark files, following python-y documentation style.
    ///
    /// For functions, they are the piece in `"""` that come right after the `def foo():` line,
    /// and they have sections for additional details. An example from a starlark file might be:
    ///
    /// ```starlark
    /// """ Module level docs here """
    ///
    /// def some_function(val: "string") -> "string":
    ///     """ This function takes a string and returns it.
    ///
    ///     This is where an explanation might go, but I have none
    ///
    ///     Args:
    ///         val: This is the value that gets returned
    ///
    ///     Returns:
    ///         The original value, because identity functions are fun.
    /// ```
    Starlark,
    /// Docstrings used with `#[starlark_module]` in rust.
    ///
    /// These are the documentation strings prefixed by `///` (like these docs) on
    /// `#[starlark_module]`, and the functions / attributes within it. It supports
    /// a section `# Arguments`, and `# Returns`, and removes some lines from code
    /// blocks that are valid for rustdoc, but not useful for people using these
    /// functions via starlark. An example might be something like:
    ///
    /// ```ignore
    /// /// These are where the module / object level docs go
    /// #[starlark_module]
    /// fn add_some_value(builder: &mut GlobalsBuilder) {
    ///     /// attr1 is an attribute that does nothing interesting.
    ///     #[attribute]
    ///     fn attr1(_this: Value<'v>) -> String {
    ///         Ok("attr1".to_owned())
    ///     }
    ///     /// Copies a string
    ///     ///
    ///     /// This is where details would be, if this were
    ///     /// a more interesting function.
    ///     ///
    ///     /// # Arguments
    ///     /// * `s`: This is string that is returned.
    ///     ///
    ///     /// # Returns
    ///     /// The a copy of the original string.
    ///     fn copy_string(s: &str) -> String {
    ///         Ok(s.to_owned())
    ///     }
    /// }
    /// ```
    Rust,
}

impl DocString {
    /// Extracts the docstring from a function or module body, iff the first
    /// statement is a string literal.
    pub(crate) fn extract_raw_starlark_docstring<P: AstPayload>(
        body: &AstStmtP<P>,
    ) -> Option<String> {
        if let StmtP::Statements(stmts) = &body.node {
            if let Some(Spanned {
                node:
                    StmtP::Expression(Spanned {
                        node: ExprP::Literal(AstLiteral::String(s)),
                        ..
                    }),
                ..
            }) = stmts.first()
            {
                return Some(s.node.to_owned());
            }
        };
        None
    }

    /// Do common work to parse a docstring (dedenting, splitting summary and details, etc)
    pub fn from_docstring(kind: DocStringKind, user_docstring: &str) -> Option<DocString> {
        let trimmed_docs = user_docstring.trim();
        if trimmed_docs.is_empty() {
            None
        } else {
            static SPLIT_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"\n[ ]*\n").unwrap());
            let split: Option<(&str, &str)> = SPLIT_RE.splitn(trimmed_docs, 2).collect_tuple();
            let (summary, details) = match split {
                Some((summary, details)) if !summary.is_empty() && !details.is_empty() => {
                    // Dedent the details separately so that people can have the summary on the
                    // same line as the opening quotes, and the details indented on subsequent
                    // lines.
                    let details = match kind {
                        DocStringKind::Starlark => textwrap::dedent(details).trim().to_owned(),
                        DocStringKind::Rust => {
                            Self::remove_rust_comments(textwrap::dedent(details).trim())
                        }
                    };
                    (summary, Some(details))
                }
                _ => (trimmed_docs, None),
            };

            // Remove any newlines (and surrounding whitespace) in the summary, and
            // replace them with a single space.
            static NEWLINES_RE: Lazy<Regex> =
                Lazy::new(|| Regex::new(r"(\S)\s*\n\s*(\S)").unwrap());
            let summary = NEWLINES_RE.replace_all(summary, r"$1 $2").to_string();

            Some(DocString { summary, details })
        }
    }

    /// Removes rustdoc-style commented out lines from code blocks.
    fn remove_rust_comments(details: &str) -> String {
        static CODEBLOCK_RE: Lazy<Regex> = Lazy::new(|| {
            RegexBuilder::new(r"```(\w*)\n.*?```")
                .dot_matches_new_line(true)
                .build()
                .expect("regex to compile")
        });
        static COMMENT_RE: Lazy<Regex> = Lazy::new(|| {
            RegexBuilder::new(r"^# .*$\n")
                .multi_line(true)
                .build()
                .expect("regex to compile")
        });
        CODEBLOCK_RE
            .replace_all(details, |caps: &regex::Captures| {
                match caps.get(1).expect("language group").as_str() {
                    "" | "rust" => COMMENT_RE
                        .replace_all(caps.get(0).expect("$0 to exist").as_str(), "")
                        .to_string(),
                    _ => caps.get(0).expect("$0 to exist").as_str().to_owned(),
                }
            })
            .to_string()
    }

    /// Join lines up, dedent them, and trim them
    fn join_and_dedent_lines(lines: &[String]) -> String {
        textwrap::dedent(&lines.join("\n")).trim().to_owned()
    }

    /// Parse the sections out of a docstring's `details` text, and remove the requested sections from the text.
    ///
    /// "sections" are the various things in doc strings like "Arguments:", "Returns:", etc
    ///
    /// # Returns
    /// - A new instance of `DocString`, with the requested sections, if found, removed.
    /// - A mapping of section name, converted to lower case, to the cleaned up section text
    ///     i.e. dedented, section header not present, etc for any found sections.
    fn parse_and_remove_sections(
        self,
        kind: DocStringKind,
        requested_sections: &[&str],
    ) -> (Self, HashMap<String, String>) {
        let mut sections = HashMap::new();

        let mut finish_section =
            |current_section: &mut Option<String>, current_section_text: &mut Vec<String>| {
                if let Some(s) = current_section.take() {
                    sections.insert(s, DocString::join_and_dedent_lines(current_section_text));
                    current_section_text.clear();
                }
            };

        static STARLARK_SECTION_RE: Lazy<Regex> =
            Lazy::new(|| Regex::new(r"^([\w -]+):\s*$").unwrap());
        static STARLARK_INDENTED_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"^(?:\s|$)").unwrap());
        static RUST_SECTION_RE: Lazy<Regex> =
            Lazy::new(|| Regex::new(r"^# ([\w -]+)\s*$").unwrap());
        static RUST_INDENTED_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"^.*").unwrap());

        let (section_re, indented_re) = match kind {
            DocStringKind::Starlark => (&STARLARK_SECTION_RE, &STARLARK_INDENTED_RE),
            DocStringKind::Rust => (&RUST_SECTION_RE, &RUST_INDENTED_RE),
        };

        if let Some(details) = self.details {
            let mut new_details = vec![];
            let mut current_section = None;
            let mut current_section_text = vec![];

            for line in details.lines() {
                if let Some(matches) = section_re.captures(line) {
                    finish_section(&mut current_section, &mut current_section_text);

                    let found_section = matches.get(1).unwrap().as_str().to_ascii_lowercase();
                    if requested_sections.contains(&found_section.as_str()) {
                        current_section = Some(found_section);
                    } else {
                        new_details.push(line.to_owned());
                    }
                } else if current_section.is_some() && indented_re.is_match(line) {
                    current_section_text.push(line.to_owned());
                } else {
                    new_details.push(line.to_owned());
                    finish_section(&mut current_section, &mut current_section_text);
                }
            }

            finish_section(&mut current_section, &mut current_section_text);

            let joined_details = new_details.join("\n").trim().to_owned();
            let details = match joined_details.is_empty() {
                true => None,
                false => Some(joined_details),
            };
            (
                Self {
                    summary: self.summary,
                    details,
                },
                sections,
            )
        } else {
            (self, sections)
        }
    }
}

/// Line / column for where in a file a symbol is.
#[derive(Debug, Clone, PartialEq, Serialize, Default)]
pub struct Pos {
    /// Line number, zero based.
    pub line: usize,
    /// Column number, zero based.
    pub column: usize,
}

/// The file a symbol resides in, and if available its location within that file.
#[derive(Debug, Clone, PartialEq, Serialize, Default)]
pub struct Location {
    /// `path` is a string that can be passed into `load()` statements.
    pub path: String,
    /// Location of the symbol within the file.
    pub position: Option<Pos>,
}

/// The main identifier for a symbol.
#[derive(Debug, Clone, PartialEq, Serialize, Default)]
pub struct Identifier {
    /// The name of the symbol (e.g. the function name, a name or path for a module, etc).
    pub name: String,
    /// Where the symbol is located, or absent if it is a built-in symbol.
    pub location: Option<Location>,
}

/// The type of a given parameter, field, etc.
#[derive(Debug, Clone, PartialEq, Serialize, Allocative)]
pub struct DocType {
    /// The type string that one would find in a starlark expression.
    pub raw_type: Ty,
}

/// Documents a full module.
#[derive(Debug, Clone, PartialEq, Serialize, Default, Allocative)]
pub struct DocModule {
    /// In general, this should be the first statement of a loaded file, if that statement is
    /// a string literal.
    pub docs: Option<DocString>,
    /// A mapping of top level symbols to their documentation, if any.
    pub members: SmallMap<String, DocMember>,
}

impl DocModule {
    fn render_as_code(&self) -> String {
        let mut res = self
            .docs
            .as_ref()
            .map(DocString::render_as_quoted_code)
            .unwrap_or_default();
        for (k, v) in &self.members {
            res.push('\n');
            res.push_str(&(Doc::named_item(k.clone(), v.clone().to_doc_item())).render_as_code());
            res.push('\n');
        }
        res
    }
}

/// Documents a single function.
#[derive(Debug, Clone, PartialEq, Serialize, Default, Allocative)]
pub struct DocFunction {
    /// Documentation for the function. If parsed, this should generally be the first statement
    /// of a function's body if that statement is a string literal. Any sections like "Args:",
    /// "Returns", etc are kept intact. It is up to the consumer to remove these sections if
    /// they are present.
    pub docs: Option<DocString>,
    /// The parameters that this function takes. Docs for these parameters should generally be
    /// extracted from the main docstring's details.
    pub params: Vec<DocParam>,
    /// Details about what this function returns.
    pub ret: DocReturn,
    /// Does this function provide a `.type` value.
    pub dot_type: Option<String>,
}

impl DocFunction {
    fn starlark_docstring(&self) -> Option<String> {
        let mut docs = String::new();
        if let Some(main_docs) = self.docs.as_ref().map(DocString::render_as_code) {
            docs.push_str(&main_docs);
        }

        let args_indentation_count = self
            .params
            .iter()
            .map(|p| match p {
                DocParam::NoArgs | DocParam::OnlyPosBefore => 0,
                DocParam::Arg { name, .. }
                | DocParam::Args { name, .. }
                | DocParam::Kwargs { name, .. } => name.len() + 2,
            })
            .max()
            .unwrap_or_default();
        let args_indentation = " ".repeat(args_indentation_count);

        let args_docs = self
            .params
            .iter()
            .filter_map(|p| p.starlark_docstring(&args_indentation))
            .join("\n");
        if !args_docs.is_empty() {
            let indented = indent_trimmed(&args_docs, "    ");
            docs.push_str(&format!("\n\nArgs:\n{}", indented));
        }

        if let Some(ret_docs) = self.ret.starlark_docstring() {
            let indented = indent_trimmed(&ret_docs, "    ");
            docs.push_str(&format!("\n\nRet:\n{}", indented));
        }
        if docs.is_empty() {
            None
        } else {
            Some(indent_trimmed(
                &format!("\"\"\"\n{}\n\"\"\"", docs.trim_start()),
                "    ",
            ))
        }
    }

    fn render_as_code(&self, name: &str) -> String {
        let params: Vec<_> = self.params.iter().map(DocParam::render_as_code).collect();
        let spacer_len = if params.is_empty() {
            0
        } else {
            (params.len() - 1) * 2
        };
        let params_len = params.iter().map(|a| a.len()).sum::<usize>() + spacer_len;
        let params = if params_len > 60 {
            format!("(\n{}\n)", indent_trimmed(&params.join(",\n"), "    "))
        } else {
            format!("({})", params.join(", "))
        };
        let docstring = self
            .starlark_docstring()
            .map(|mut ds| {
                ds.push('\n');
                ds
            })
            .unwrap_or_default();
        let ret = self
            .ret
            .typ
            .as_ref()
            .map(|t| format!(" -> {}", t.raw_type))
            .unwrap_or_default();

        format!("def {}{}{}:\n{}    pass", name, params, ret, docstring)
    }

    /// Parses function documentation out of a docstring
    ///
    /// # Arguments
    /// * `kind`: The kind of docstring. This determines the formatting that is parsed.
    /// * `params`: The parameters of the function.
    /// * `return_type`: The return type. This is pulled from typing info / directly from users,
    ///                  so it cannot be inferred generically.
    /// * `raw_docstring`: The raw docstring to be parsed and potentially modified,
    ///                    removing the sections detailing arguments and return values.
    ///                    The format is determined by `kind`.
    pub fn from_docstring(
        kind: DocStringKind,
        mut params: Vec<DocParam>,
        return_type: Option<DocType>,
        raw_docstring: Option<&str>,
        dot_type: Option<String>,
    ) -> Self {
        match raw_docstring.and_then(|raw| DocString::from_docstring(kind, raw)) {
            Some(ds) => {
                let (function_docstring, sections) =
                    ds.parse_and_remove_sections(kind, &["arguments", "args", "returns", "return"]);

                match sections.get("arguments").or_else(|| sections.get("args")) {
                    Some(args) => {
                        let entries = Self::parse_params(kind, args);
                        for x in &mut params {
                            match x {
                                DocParam::Arg { name, docs, .. }
                                | DocParam::Args { name, docs, .. }
                                | DocParam::Kwargs { name, docs, .. } => match entries.get(name) {
                                    Some(raw) => *docs = DocString::from_docstring(kind, raw),
                                    _ => (),
                                },
                                _ => (),
                            }
                        }
                    }
                    _ => (),
                }

                let return_docs = sections
                    .get("return")
                    .or_else(|| sections.get("returns"))
                    .and_then(|raw| DocString::from_docstring(kind, raw));

                DocFunction {
                    docs: Some(function_docstring),
                    params,
                    ret: DocReturn {
                        docs: return_docs,
                        typ: return_type,
                    },
                    dot_type,
                }
            }
            None => DocFunction {
                docs: None,
                params,
                ret: DocReturn {
                    docs: None,
                    typ: return_type,
                },
                dot_type,
            },
        }
    }

    /// Parse out parameter docs from an "Args:" section of a docstring
    ///
    /// `args_section` should be dedented, and generally should just be the `args` key of
    /// the `DocString::parse_and_remove_sections()` function call. This is done as a
    /// separate function to reduce the number of times that sections are parsed out of
    /// docstring (e.g. if a user wants both the `Args:` and `Returns:` sections)
    fn parse_params(kind: DocStringKind, args_section: &str) -> HashMap<String, String> {
        static STARLARK_ARG_RE: Lazy<Regex> =
            Lazy::new(|| Regex::new(r"^(\*{0,2}\w+):\s*(.*)").unwrap());
        static RUST_ARG_RE: Lazy<Regex> =
            Lazy::new(|| Regex::new(r"^(?:\* )?`(\w+)`:?\s*(.*)").unwrap());

        static INDENTED_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"^(?:\s|$)").unwrap());

        let arg_re = match kind {
            DocStringKind::Starlark => &STARLARK_ARG_RE,
            DocStringKind::Rust => &RUST_ARG_RE,
        };

        let mut ret = HashMap::new();
        let mut current_arg = None;
        let mut current_text = vec![];

        for line in args_section.lines() {
            if let Some(matches) = arg_re.captures(line) {
                if let Some(a) = current_arg.take() {
                    ret.insert(a, DocString::join_and_dedent_lines(&current_text));
                }

                current_arg = Some(matches.get(1).unwrap().as_str().to_owned());

                let doc_match = matches.get(2).unwrap();
                current_text = vec![format!(
                    "{}{}",
                    " ".repeat(doc_match.start()),
                    doc_match.as_str()
                )];
            } else if current_arg.is_some() && INDENTED_RE.is_match(line) {
                current_text.push(line.to_owned());
            }
        }

        if let Some(a) = current_arg.take() {
            ret.insert(a, DocString::join_and_dedent_lines(&current_text));
        }

        ret
    }
}

/// A single parameter of a function.
#[derive(Debug, Clone, PartialEq, Serialize, Allocative)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum DocParam {
    /// A regular parameter that may or may not have a default value.
    Arg {
        name: String,
        docs: Option<DocString>,
        #[serde(rename = "type")]
        typ: Option<DocType>,
        /// If present, this parameter has a default value. This is the `repr()` of that value.
        default_value: Option<String>,
    },
    /// Represents the "*" argument.
    NoArgs,
    /// Represents the "/" argument from [PEP 570](https://peps.python.org/pep-0570/).
    OnlyPosBefore,
    /// Represents the "*args" style of argument.
    Args {
        name: String,
        docs: Option<DocString>,
        #[serde(rename = "type")]
        typ: Option<DocType>,
    },
    /// Represents the "**kwargs" style of argument.
    Kwargs {
        name: String,
        docs: Option<DocString>,
        #[serde(rename = "type")]
        typ: Option<DocType>,
    },
}

impl DocParam {
    fn starlark_docstring(&self, max_indentation: &str) -> Option<String> {
        let (name, docs) = match self {
            DocParam::Arg { name, docs, .. } => Some((name, docs)),
            DocParam::NoArgs | DocParam::OnlyPosBefore => None,
            DocParam::Args { name, docs, .. } => Some((name, docs)),
            DocParam::Kwargs { name, docs, .. } => Some((name, docs)),
        }?;
        let rendered_docs = docs.as_ref()?.render_as_code();
        let mut indented = indent_trimmed(&rendered_docs, max_indentation);
        indented.replace_range(..name.len() + 2, &format!("{}: ", name));
        Some(indented)
    }

    fn render_as_code(&self) -> String {
        match self {
            DocParam::Arg {
                name,
                typ,
                default_value,
                ..
            } => match (typ.as_ref(), default_value.as_ref()) {
                (Some(t), Some(default)) => format!("{}: {} = {}", name, t.raw_type, default),
                (Some(t), None) => format!("{}: {}", name, t.raw_type),
                (None, Some(default)) => format!("{} = {}", name, default),
                (None, None) => name.clone(),
            },
            DocParam::NoArgs => "*".to_owned(),
            DocParam::OnlyPosBefore => "/".to_owned(),
            DocParam::Args { name, typ, .. } | DocParam::Kwargs { name, typ, .. } => {
                match typ.as_ref() {
                    Some(typ) => format!("{}: {}", name, typ.raw_type),
                    None => name.clone(),
                }
            }
        }
    }
}

/// Details about the return value of a function.
#[derive(Debug, Clone, PartialEq, Serialize, Default, Allocative)]
pub struct DocReturn {
    /// Extra semantic details around the returned value's meaning.
    pub docs: Option<DocString>,
    #[serde(rename = "type")]
    pub typ: Option<DocType>,
}

impl DocReturn {
    fn starlark_docstring(&self) -> Option<String> {
        self.docs.as_ref().map(DocString::render_as_code)
    }
}

/// A single property of an object. These are explicitly not functions (see [`DocMember`]).
#[derive(Debug, Clone, PartialEq, Serialize, Default, Allocative)]
pub struct DocProperty {
    pub docs: Option<DocString>,
    #[serde(rename = "type")]
    pub typ: Option<DocType>,
}

impl DocProperty {
    fn render_as_code(&self, name: &str) -> String {
        match (
            self.typ.as_ref(),
            self.docs.as_ref().map(DocString::render_as_quoted_code),
        ) {
            // TODO(nmj): The starlark syntax needs to be updated to support type
            //            annotations on values as python does. Afterward, use these
            //            format strings.
            // (Some(t), Some(ds)) => {
            //     format!("{}\n_{}: {} = None", ds, name, t.raw_type)
            // }
            // (Some(t), None) => format!(r#"_{}: {} = None"#, name, t.raw_type),
            (Some(t), Some(ds)) => {
                format!("{}\n# type: {}\n_{} = None", ds, t.raw_type, name)
            }
            (Some(t), None) => format!("# type: {}\n_{} = None", t.raw_type, name),

            (None, Some(ds)) => format!("{}\n_{} = None", ds, name),
            (None, None) => format!("_{} = None", name),
        }
    }
}

/// A named member of an object.
#[derive(Debug, Clone, PartialEq, Serialize, Allocative)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum DocMember {
    Property(DocProperty),
    Function(DocFunction),
}

impl DocMember {
    pub(crate) fn from_value(value: Value) -> Self {
        // If we have a value which is a complex type, the right type to put in the docs is not the type
        // it represents, but it's just a property we should point at
        match value.documentation() {
            Some(DocItem::Function(x)) => DocMember::Function(x),
            Some(DocItem::Property(x)) => DocMember::Property(x),
            _ => DocMember::Property(DocProperty {
                docs: None,
                typ: Some(DocType {
                    raw_type: value.get_type_starlark_repr(),
                }),
            }),
        }
    }

    pub fn to_doc_item(self) -> DocItem {
        match self {
            DocMember::Property(x) => DocItem::Property(x),
            DocMember::Function(x) => DocItem::Function(x),
        }
    }
}

/// An object with named functions/properties.
#[derive(Debug, Clone, PartialEq, Serialize, Default, Allocative)]
pub struct DocObject {
    pub docs: Option<DocString>,
    /// Name and details of each member of this object.
    pub members: SmallMap<String, DocMember>,
}

impl DocObject {
    fn render_as_code(&self, name: &str) -> String {
        let summary = self
            .docs
            .as_ref()
            .map(|ds| {
                let mut s = ds.render_as_quoted_code();
                s.push('\n');
                s
            })
            .unwrap_or_default();

        let member_docs = self
            .members
            .iter()
            .map(|(name, member)| match member {
                DocMember::Property(p) => p.render_as_code(name),
                DocMember::Function(f) => f.render_as_code(&format!("_{}", name)),
            })
            .join("\n\n");

        let exported_struct_members = self
            .members
            .iter()
            .map(|(name, _)| format!("    {} = _{},", name, name))
            .join("\n");
        let exported_struct = if !exported_struct_members.is_empty() {
            format!(
                "{}{} = struct(\n{}\n)",
                summary, name, exported_struct_members
            )
        } else {
            String::new()
        };

        format!("{}\n\n{}", member_docs, exported_struct)
            .trim()
            .to_owned()
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Allocative)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum DocItem {
    Module(DocModule),
    Object(DocObject),
    Function(DocFunction),
    Property(DocProperty),
}

/// The main structure that represents the documentation for a given symbol / module.
#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct Doc {
    pub id: Identifier,
    pub item: DocItem,
    /// Custom key-value pairs that are not interpreted directly by starlark, and can be
    /// used as arbitrary data for documentation tooling.
    pub custom_attrs: HashMap<String, String>,
}

impl Doc {
    pub fn named_item(name: String, item: DocItem) -> Self {
        Doc {
            id: Identifier {
                name,
                location: None,
            },
            item,
            custom_attrs: HashMap::new(),
        }
    }

    /// Render a starlark code representation of this documentation object.
    ///
    /// Function bodies for these consist of a single "pass" statement, and objects
    /// are represented as structs.
    pub fn render_as_code(&self) -> String {
        match &self.item {
            DocItem::Module(m) => m.render_as_code(),
            DocItem::Object(o) => o.render_as_code(&self.id.name),
            DocItem::Function(f) => f.render_as_code(&self.id.name),
            DocItem::Property(p) => p.render_as_code(&self.id.name),
        }
    }
}

/// Render a series of [`Doc`] objects into a "starlark" file.
///
/// Function bodies for these consist of a single "pass" statement, and objects
/// are represented as structs.
///
/// The returned array may not be in the same order as the originally provided docs.
/// They are in the order that they should appear in the rendered starlark file.
pub fn render_docs_as_code(docs: &[Doc]) -> String {
    let (modules, non_modules): (Vec<_>, Vec<_>) = docs
        .iter()
        .partition(|d| matches!(d.item, DocItem::Module(_)));
    modules
        .into_iter()
        .chain(non_modules.into_iter())
        .map(|d| d.render_as_code())
        .join("\n\n")
}

/// Get documentation for all items registered with `#[derive(StarlarkDocs)]`
///
/// Note: Because `StarlarkDocs` uses the inventory crate under the hood, in statically linked
/// binaries, documentation from all compiled crates in the binary will be included.
///
/// For dynamically linked binaries, documentation will only be able to retrieved after the crate's
/// library is `dlopen()`ed.
pub fn get_registered_starlark_docs() -> Vec<Doc> {
    inventory::iter::<RegisteredDoc>
        .into_iter()
        .filter_map(|d| (d.getter)())
        .collect()
}

#[doc(hidden)]
pub struct RegisteredDoc {
    pub getter: fn() -> Option<Doc>,
}

inventory::collect!(RegisteredDoc);

impl RegisteredDoc {
    /// This function is called from generated code.
    pub fn for_type<'v, T: StarlarkValue<'v>>(custom_attrs: &[(&str, &str)]) -> Option<Doc> {
        let name = T::TYPE.to_owned();
        let id = Identifier {
            name,
            location: None,
        };
        let item = DocItem::Object(T::get_methods()?.documentation());
        let custom_attrs = custom_attrs
            .iter()
            .map(|(k, v)| ((*k).to_owned(), (*v).to_owned()))
            .collect();
        Some(Doc {
            id,
            item,
            custom_attrs,
        })
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn parses_starlark_docstring() {
        assert_eq!(
            DocString::from_docstring(DocStringKind::Starlark, " "),
            None
        );
        assert_eq!(
            DocString::from_docstring(
                DocStringKind::Starlark,
                " \n\nThis should be the summary\n\n"
            ),
            Some(DocString {
                summary: "This should be the summary".to_owned(),
                details: None,
            })
        );
        assert_eq!(
            DocString::from_docstring(
                DocStringKind::Starlark,
                " \n\nThis should be the summary\n\n "
            ),
            Some(DocString {
                summary: "This should be the summary".to_owned(),
                details: None,
            })
        );
        assert_eq!(
            DocString::from_docstring(
                DocStringKind::Starlark,
                "Summary line here\n    \nDetails after some spaces\n\nand some more newlines"
            ),
            Some(DocString {
                summary: "Summary line here".to_owned(),
                details: Some("Details after some spaces\n\nand some more newlines".to_owned()),
            })
        );
        assert_eq!(
            DocString::from_docstring(
                DocStringKind::Starlark,
                r#"
        This is the summary.
          It has multiple lines and some spaces, and should be collapsed

        This should be a multiline set of details.
        It should be:
            - Dedented
            - Trimmed
            - Split properly from the summary

"#
            ),
            Some(DocString {
                summary: "This is the summary. It has multiple lines and some spaces, and should be collapsed".to_owned(),
                details: Some(
                    concat!(
                        "This should be a multiline set of details.\n",
                        "It should be:\n",
                        "    - Dedented\n",
                        "    - Trimmed\n",
                        "    - Split properly from the summary"
                    )
                    .to_owned()
                ),
            })
        );
        assert_eq!(
            DocString::from_docstring(
                DocStringKind::Starlark,
                r#"This is a summary line that is not dedented like the 'details'

        Typing the first line right after the """ in python docstrings is common,
        while putting the rest of the docstring indented. Just support both so it
        doesn't surprise anyone.
        "#
            ),
            Some(DocString {
                summary: "This is a summary line that is not dedented like the 'details'"
                    .to_owned(),
                details: Some(
                    concat!(
                "Typing the first line right after the \"\"\" in python docstrings is common,\n",
                "while putting the rest of the docstring indented. Just support both so it\n",
                "doesn't surprise anyone."
                )
                    .to_owned()
                ),
            })
        );
    }

    #[test]
    fn parses_rust_docstring() {
        let raw = r#"
        This is the summary line
          that sometimes is split on two lines

        This is the second part. It has some code blocks

        ```
        # foo() {
        "bar"
        # }
        ```

        ```python
        # This is a python comment. Leave it be
        print(1)
        ```

        ```rust
        # other_foo() {
        "other_bar"
        # }
        ```
        "#;

        let parsed = DocString::from_docstring(DocStringKind::Rust, raw).unwrap();
        assert_eq!(
            "This is the summary line that sometimes is split on two lines",
            parsed.summary
        );
        assert_eq!(
            concat!(
                "This is the second part. It has some code blocks\n\n",
                "```\n",
                "\"bar\"\n",
                "```\n\n",
                "```python\n",
                "# This is a python comment. Leave it be\n",
                "print(1)\n",
                "```\n\n",
                "```rust\n",
                "\"other_bar\"",
                "\n```"
            ),
            parsed.details.unwrap()
        );
    }

    #[test]
    fn parses_and_removes_sections_from_starlark_docstring() {
        let raw_docs = r#"This is an example docstring

        We have some details up here that should not be parsed

        Some empty section:
        Example:
            First line of the section

            A newline with no space after it before the second one,
                and a third that's indented further.
        This is not in the example section

        Last:
            This is something in the last section
        "#;
        let expected_docstring = DocString::from_docstring(
            DocStringKind::Starlark,
            r#"This is an example docstring

            We have some details up here that should not be parsed

            Some empty section:
            This is not in the example section

            Last:
                This is something in the last section
            "#,
        )
        .unwrap();

        let expected_sections = HashMap::from([(
            "example".to_owned(),
            concat!(
                "First line of the section\n\n",
                "A newline with no space after it before the second one,\n",
                "    and a third that's indented further."
            )
            .to_owned(),
        )]);

        let ds = DocString::from_docstring(DocStringKind::Starlark, raw_docs).unwrap();
        let (new_ds, sections) =
            ds.parse_and_remove_sections(DocStringKind::Starlark, &["example"]);

        assert_eq!(new_ds, expected_docstring);
        assert_eq!(sections, expected_sections);
    }

    #[test]
    fn parses_and_removes_sections_from_rust_docstring() {
        let raw_docs = r#"This is an example docstring

        We have some details up here that should not be parsed

        # Some Section

        ```
        # This is a commented out line in a codeblock
        fn some_func() {}
        ```

        # Example
        First line of the section

        Note that, unlike starlark doc strings,
        we don't require indentation. The end of a
        section is either a new section appearing,
        or the end of the string.

        # Last
        This is something in the last section
        "#;
        let expected_docstring = DocString::from_docstring(
            DocStringKind::Rust,
            r#"This is an example docstring

        We have some details up here that should not be parsed

        # Some Section

        ```
        fn some_func() {}
        ```

        # Last
        This is something in the last section
        "#,
        )
        .unwrap();

        let expected_sections = HashMap::from([(
            "example".to_owned(),
            concat!(
                "First line of the section\n\n",
                "Note that, unlike starlark doc strings,\n",
                "we don't require indentation. The end of a\n",
                "section is either a new section appearing,\n",
                "or the end of the string.",
            )
            .to_owned(),
        )]);

        let ds = DocString::from_docstring(DocStringKind::Rust, raw_docs).unwrap();
        let (new_ds, sections) = ds.parse_and_remove_sections(DocStringKind::Rust, &["example"]);

        assert_eq!(new_ds, expected_docstring);
        assert_eq!(sections, expected_sections);
    }

    fn arg(name: &str) -> DocParam {
        DocParam::Arg {
            name: name.to_owned(),
            docs: None,
            typ: None,
            default_value: None,
        }
    }

    #[test]
    fn parses_starlark_function_docstring() {
        let docstring = r#"This is an example docstring

        Details here

        Args:
            arg_foo: The argument named foo
            arg_bar: The argument named bar. It has
                     a longer doc string that spans
                     over three lines
            *args: Docs for args
            **kwargs: Docs for kwargs

        Returns:
            A value
        "#;

        let kind = DocStringKind::Starlark;
        let return_type = Some(DocType {
            raw_type: Ty::int(),
        });
        let expected = DocFunction {
            docs: DocString::from_docstring(kind, "This is an example docstring\n\nDetails here"),
            params: vec![
                DocParam::Arg {
                    name: "**kwargs".to_owned(),
                    docs: DocString::from_docstring(kind, "Docs for kwargs"),
                    typ: None,
                    default_value: None,
                },
                DocParam::Arg {
                    name: "*args".to_owned(),
                    docs: DocString::from_docstring(kind, "Docs for args"),
                    typ: None,
                    default_value: None,
                },
                DocParam::Arg {
                    name: "arg_bar".to_owned(),
                    docs: DocString::from_docstring(
                        kind,
                        concat!(
                            "The argument named bar. It has\n",
                            "a longer doc string that spans\n",
                            "over three lines"
                        ),
                    ),
                    typ: None,
                    default_value: None,
                },
                DocParam::Arg {
                    name: "arg_foo".to_owned(),
                    docs: DocString::from_docstring(kind, "The argument named foo"),
                    typ: None,
                    default_value: None,
                },
            ],
            ret: DocReturn {
                docs: DocString::from_docstring(kind, "A value"),
                typ: return_type.clone(),
            },
            dot_type: None,
        };

        let function_docs = DocFunction::from_docstring(
            kind,
            vec![
                arg("**kwargs"),
                arg("*args"),
                arg("arg_bar"),
                arg("arg_foo"),
            ],
            return_type,
            Some(docstring),
            None,
        );

        assert_eq!(expected, function_docs);
    }

    #[test]
    fn parses_rust_function_docstring() {
        let docstring = r#"This is an example docstring

        Details here

        # Arguments
        * `arg_foo`: The argument named foo
        `arg_bar`: The argument named bar. It has
                   a longer doc string that spans
                   over three lines

        # Returns
        A value
        "#;

        let kind = DocStringKind::Rust;
        let return_type = Some(DocType {
            raw_type: Ty::int(),
        });
        let expected = DocFunction {
            docs: DocString::from_docstring(kind, "This is an example docstring\n\nDetails here"),
            params: vec![
                DocParam::Arg {
                    name: "arg_bar".to_owned(),
                    docs: DocString::from_docstring(
                        kind,
                        concat!(
                            "The argument named bar. It has\n",
                            "a longer doc string that spans\n",
                            "over three lines"
                        ),
                    ),
                    typ: None,
                    default_value: None,
                },
                DocParam::Arg {
                    name: "arg_foo".to_owned(),
                    docs: DocString::from_docstring(kind, "The argument named foo"),
                    typ: None,
                    default_value: None,
                },
            ],
            ret: DocReturn {
                docs: DocString::from_docstring(kind, "A value"),
                typ: return_type.clone(),
            },
            dot_type: None,
        };

        let function_docs = DocFunction::from_docstring(
            kind,
            vec![arg("arg_bar"), arg("arg_foo")],
            return_type,
            Some(docstring),
            None,
        );

        assert_eq!(expected, function_docs);
    }
}


/*
 * 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::sync::Arc;

use allocative::Allocative;
use dupe::Dupe;
use itertools::Itertools;
use once_cell::sync::Lazy;
use once_cell::sync::OnceCell;

use crate::collections::symbol_map::Symbol;
use crate::collections::symbol_map::SymbolMap;
use crate::collections::Hashed;
use crate::collections::SmallMap;
use crate::docs::DocMember;
use crate::docs::DocModule;
use crate::docs::DocObject;
use crate::docs::DocString;
use crate::docs::DocStringKind;
use crate::stdlib;
pub use crate::stdlib::LibraryExtension;
use crate::typing::Ty;
use crate::values::function::NativeAttribute;
use crate::values::function::NativeCallableRawDocs;
use crate::values::function::NativeFunc;
use crate::values::function::NativeMeth;
use crate::values::function::SpecialBuiltinFunction;
use crate::values::layout::value_not_special::FrozenValueNotSpecial;
use crate::values::structs::AllocStruct;
use crate::values::types::function::NativeFunction;
use crate::values::types::function::NativeMethod;
use crate::values::AllocFrozenValue;
use crate::values::FrozenHeap;
use crate::values::FrozenHeapRef;
use crate::values::FrozenStringValue;
use crate::values::FrozenValue;
use crate::values::Heap;
use crate::values::Value;

/// The global values available during execution.
#[derive(Clone, Dupe, Debug, Allocative)]
pub struct Globals(Arc<GlobalsData>);

/// Methods of an object.
#[derive(Clone, Debug)]
pub struct Methods {
    /// This field holds the objects referenced in `members`.
    #[allow(dead_code)]
    heap: FrozenHeapRef,
    members: SymbolMap<FrozenValueNotSpecial>,
    docstring: Option<String>,
}

#[derive(Debug, Allocative)]
struct GlobalsData {
    heap: FrozenHeapRef,
    variables: SymbolMap<FrozenValue>,
    variable_names: Vec<FrozenStringValue>,
    docstring: Option<String>,
}

/// Used to build a [`Globals`] value.
#[derive(Debug)]
pub struct GlobalsBuilder {
    // The heap everything is allocated in
    heap: FrozenHeap,
    // Normal top-level variables, e.g. True/hash
    variables: SymbolMap<FrozenValue>,
    // The list of struct fields, pushed to the end
    struct_fields: Vec<SmallMap<FrozenStringValue, FrozenValue>>,
    // The raw docstring for this module
    docstring: Option<String>,
}

/// Used to build a [`Methods`] value.
#[derive(Debug)]
pub struct MethodsBuilder {
    /// The heap everything is allocated in.
    heap: FrozenHeap,
    /// Members, either `NativeMethod` or `NativeAttribute`.
    members: SymbolMap<FrozenValueNotSpecial>,
    /// The raw docstring for the main object.
    docstring: Option<String>,
}

impl Globals {
    /// Create an empty [`Globals`], with no functions in scope.
    pub fn new() -> Self {
        GlobalsBuilder::new().build()
    }

    /// Create a [`Globals`] following the
    /// [Starlark standard](https://github.com/bazelbuild/starlark/blob/master/spec.md#built-in-constants-and-functions).
    pub fn standard() -> Self {
        GlobalsBuilder::standard().build()
    }

    /// Create a [`Globals`] combining those functions in the Starlark standard plus
    /// all those defined in [`LibraryExtension`].
    pub fn extended() -> Self {
        GlobalsBuilder::extended().build()
    }

    /// Empty globals.
    pub(crate) fn empty() -> &'static Globals {
        static EMPTY: Lazy<Globals> = Lazy::new(|| GlobalsBuilder::new().build());
        &EMPTY
    }

    /// Create a [`Globals`] combining those functions in the Starlark standard plus
    /// all those given in the [`LibraryExtension`] arguments.
    pub fn extended_by(extensions: &[LibraryExtension]) -> Self {
        GlobalsBuilder::extended_by(extensions).build()
    }

    /// This function is only safe if you first call `heap` and keep a reference to it.
    /// Therefore, don't expose it on the public API.
    pub(crate) fn get<'v>(&'v self, name: &str) -> Option<Value<'v>> {
        self.get_frozen(name).map(FrozenValue::to_value)
    }

    /// This function is only safe if you first call `heap` and keep a reference to it.
    /// Therefore, don't expose it on the public API.
    pub(crate) fn get_frozen(&self, name: &str) -> Option<FrozenValue> {
        self.0.variables.get_str(name).copied()
    }

    /// Get all the names defined in this environment.
    pub fn names(&self) -> impl Iterator<Item = FrozenStringValue> + '_ {
        self.0.variable_names.iter().copied()
    }

    /// Iterate over all the items in this environment.
    /// Note returned values are owned by this globals.
    pub(crate) fn iter(&self) -> impl Iterator<Item = (&str, FrozenValue)> {
        self.0.variables.iter().map(|(n, v)| (n.as_str(), *v))
    }

    pub(crate) fn heap(&self) -> &FrozenHeapRef {
        &self.0.heap
    }

    /// Print information about the values in this object.
    pub fn describe(&self) -> String {
        self.0
            .variables
            .iter()
            .map(|(name, val)| val.to_value().describe(name.as_str()))
            .join("\n")
    }

    /// Get the documentation for both the object itself, and its members.
    pub fn documentation(&self) -> DocModule {
        let DocObject { docs, members } = common_documentation(
            &self.0.docstring,
            self.0.variables.iter().map(|(n, v)| (n.as_str(), *v)),
        );
        DocModule { docs, members }
    }
}

impl Methods {
    pub(crate) fn get<'v>(&'v self, name: &str) -> Option<Value<'v>> {
        self.get_frozen(name).map(FrozenValueNotSpecial::to_value)
    }

    pub(crate) fn get_frozen(&self, name: &str) -> Option<FrozenValueNotSpecial> {
        self.members.get_str(name).copied()
    }

    pub(crate) fn get_hashed(&self, name: Hashed<&str>) -> Option<FrozenValueNotSpecial> {
        self.members.get_hashed_str(name).copied()
    }

    pub(crate) fn get_frozen_symbol(&self, name: &Symbol) -> Option<FrozenValueNotSpecial> {
        self.members.get(name).copied()
    }

    pub(crate) fn names(&self) -> Vec<String> {
        self.members.keys().map(|x| x.as_str().to_owned()).collect()
    }

    pub(crate) fn members(&self) -> impl Iterator<Item = (&str, FrozenValue)> {
        self.members
            .iter()
            .map(|(k, v)| (k.as_str(), v.to_frozen_value()))
    }

    /// Fetch the documentation.
    pub fn documentation(&self) -> DocObject {
        common_documentation(
            &self.docstring,
            self.members
                .iter()
                .map(|(n, v)| (n.as_str(), v.to_frozen_value())),
        )
    }
}

impl GlobalsBuilder {
    /// Create an empty [`GlobalsBuilder`], with no functions in scope.
    pub fn new() -> Self {
        Self {
            heap: FrozenHeap::new(),
            variables: SymbolMap::new(),
            struct_fields: Vec::new(),
            docstring: None,
        }
    }

    /// Create a [`GlobalsBuilder`] following the
    /// [Starlark standard](https://github.com/bazelbuild/starlark/blob/master/spec.md#built-in-constants-and-functions).
    pub fn standard() -> Self {
        stdlib::standard_environment()
    }

    /// Create a [`GlobalsBuilder`] combining those functions in the Starlark standard plus
    /// all those defined in [`LibraryExtension`].
    pub fn extended() -> Self {
        Self::extended_by(LibraryExtension::all())
    }

    /// Create a [`GlobalsBuilder`] combining those functions in the Starlark standard plus
    /// all those defined in [`LibraryExtension`].
    pub fn extended_by(extensions: &[LibraryExtension]) -> Self {
        let mut res = Self::standard();
        for x in extensions {
            x.add(&mut res);
        }
        res
    }

    /// Add a nested struct to the builder. If `f` adds the definition `foo`,
    /// it will end up on a struct `name`, accessible as `name.foo`.
    pub fn struct_(&mut self, name: &str, f: impl FnOnce(&mut GlobalsBuilder)) {
        self.struct_fields.push(SmallMap::new());
        f(self);
        let fields = self.struct_fields.pop().unwrap();
        self.set(name, AllocStruct(fields));
    }

    /// A fluent API for modifying [`GlobalsBuilder`] and returning the result.
    pub fn with(mut self, f: impl FnOnce(&mut Self)) -> Self {
        f(&mut self);
        self
    }

    /// A fluent API for modifying [`GlobalsBuilder`] using [`struct_`](GlobalsBuilder::struct_).
    pub fn with_struct(mut self, name: &str, f: impl Fn(&mut GlobalsBuilder)) -> Self {
        self.struct_(name, f);
        self
    }

    /// Called at the end to build a [`Globals`].
    pub fn build(self) -> Globals {
        let variable_names = self
            .variables
            .keys()
            .map(|x| self.heap.alloc_str_intern(x.as_str()))
            .collect();
        Globals(Arc::new(GlobalsData {
            heap: self.heap.into_ref(),
            variables: self.variables,
            variable_names,
            docstring: self.docstring,
        }))
    }

    /// Set a value in the [`GlobalsBuilder`].
    pub fn set<'v, V: AllocFrozenValue>(&'v mut self, name: &str, value: V) {
        let value = value.alloc_frozen_value(&self.heap);
        match self.struct_fields.last_mut() {
            None => self.variables.insert(name, value),
            Some(fields) => {
                let name = self.heap.alloc_str(name);
                fields.insert(name, value)
            }
        };
    }

    /// Set a method. This function is usually called from code
    /// generated by `starlark_derive` and rarely needs to be called manually.
    pub fn set_function<F>(
        &mut self,
        name: &str,
        speculative_exec_safe: bool,
        raw_docs: NativeCallableRawDocs,
        typ: Option<FrozenStringValue>,
        ty: Option<Ty>,
        special_builtin_function: Option<SpecialBuiltinFunction>,
        f: F,
    ) where
        F: NativeFunc,
    {
        self.set(
            name,
            NativeFunction {
                function: Box::new(f),
                name: name.to_owned(),
                speculative_exec_safe,
                typ,
                ty,
                raw_docs: Some(raw_docs),
                special_builtin_function,
            },
        )
    }

    /// Heap where globals are allocated. Can be used to allocate additional values.
    pub fn frozen_heap(&self) -> &FrozenHeap {
        &self.heap
    }

    /// Allocate a value using the same underlying heap as the [`GlobalsBuilder`],
    /// only intended for values that are referred to by those which are passed
    /// to [`set`](GlobalsBuilder::set).
    pub fn alloc<'v, V: AllocFrozenValue>(&'v self, value: V) -> FrozenValue {
        value.alloc_frozen_value(&self.heap)
    }

    /// Set per module docstring.
    ///
    /// This function is called by the `starlark_derive` generated code
    /// and rarely needs to be called manually.
    pub fn set_docstring(&mut self, docstring: &str) {
        self.docstring = Some(docstring.to_owned());
    }
}

impl Methods {
    /// Create an empty [`Globals`], with no functions in scope.
    pub fn new() -> Self {
        MethodsBuilder::new().build()
    }
}

impl MethodsBuilder {
    /// Create an empty [`MethodsBuilder`], with no functions in scope.
    pub fn new() -> Self {
        MethodsBuilder {
            heap: FrozenHeap::new(),
            members: SymbolMap::new(),
            docstring: None,
        }
    }

    /// Called at the end to build a [`Methods`].
    pub fn build(self) -> Methods {
        Methods {
            heap: self.heap.into_ref(),
            members: self.members,
            docstring: self.docstring,
        }
    }

    /// A fluent API for modifying [`MethodsBuilder`] and returning the result.
    pub fn with(mut self, f: impl FnOnce(&mut Self)) -> Self {
        f(&mut self);
        self
    }

    /// Set the raw docstring for this object.
    pub fn set_docstring(&mut self, docstring: &str) {
        self.docstring = Some(docstring.to_owned());
    }

    /// Set a constant value in the [`MethodsBuilder`] that will be suitable for use with
    /// [`StarlarkValue::get_methods`](crate::values::StarlarkValue::get_methods).
    pub fn set_attribute<'v, V: AllocFrozenValue>(
        &'v mut self,
        name: &str,
        value: V,
        docstring: Option<String>,
    ) {
        // We want to build an attribute, that ignores its self argument, and does no subsequent allocation.
        let value = self.heap.alloc(value);
        self.set_attribute_fn(
            name,
            true,
            docstring,
            Ty::name(value.to_value().get_type()),
            move |_, _| Ok(value.to_value()),
        );
    }

    /// Set an attribute. This function is usually called from code
    /// generated by `starlark_derive` and rarely needs to be called manually.
    pub fn set_attribute_fn<F>(
        &mut self,
        name: &str,
        speculative_exec_safe: bool,
        docstring: Option<String>,
        typ: Ty,
        f: F,
    ) where
        F: for<'v> Fn(Value<'v>, &'v Heap) -> anyhow::Result<Value<'v>> + Send + Sync + 'static,
    {
        self.members.insert(
            name,
            FrozenValueNotSpecial::new(self.heap.alloc(NativeAttribute {
                function: Box::new(f),
                speculative_exec_safe,
                docstring,
                typ,
            }))
            .unwrap(),
        );
    }

    /// Set a method. This function is usually called from code
    /// generated by `starlark_derive` and rarely needs to be called manually.
    pub fn set_method<F>(
        &mut self,
        name: &str,
        speculative_exec_safe: bool,
        raw_docs: NativeCallableRawDocs,
        typ: Option<FrozenValue>,
        f: F,
    ) where
        F: NativeMeth,
    {
        self.members.insert(
            name,
            FrozenValueNotSpecial::new(self.heap.alloc(NativeMethod {
                function: Box::new(f),
                name: name.to_owned(),
                typ,
                speculative_exec_safe,
                raw_docs,
            }))
            .unwrap(),
        );
    }

    /// Allocate a value using the same underlying heap as the [`MethodsBuilder`]
    pub fn alloc<'v, V: AllocFrozenValue>(&'v self, value: V) -> FrozenValue {
        value.alloc_frozen_value(&self.heap)
    }
}

/// Used to create methods for a [`StarlarkValue`](crate::values::StarlarkValue).
///
/// To define a method `foo()` on your type, define
///  usually written as:
///
/// ```ignore
/// fn my_methods(builder: &mut GlobalsBuilder) {
///     fn foo(me: ARef<Foo>) -> anyhow::Result<NoneType> {
///         ...
///     }
/// }
///
/// impl StarlarkValue<'_> for Foo {
///     ...
///     fn get_methods(&self) -> Option<&'static Globals> {
///         static RES: GlobalsStatic = GlobalsStatic::new();
///         RES.methods(module_creator)
///     }
///     ...
/// }
/// ```
pub struct GlobalsStatic(OnceCell<Globals>);

/// Similar to [`GlobalsStatic`], but for methods.
pub struct MethodsStatic(OnceCell<Methods>);

impl GlobalsStatic {
    /// Create a new [`GlobalsStatic`].
    pub const fn new() -> Self {
        Self(OnceCell::new())
    }

    fn globals(&'static self, x: impl FnOnce(&mut GlobalsBuilder)) -> &'static Globals {
        self.0.get_or_init(|| GlobalsBuilder::new().with(x).build())
    }

    /// Get a function out of the object. Requires that the function passed only set a single
    /// value. If populated via a `#[starlark_module]`, that means a single function in it.
    pub fn function(&'static self, x: impl FnOnce(&mut GlobalsBuilder)) -> FrozenValue {
        let globals = self.globals(x);
        assert!(
            globals.0.variables.len() == 1,
            "GlobalsBuilder.function must have exactly 1 member, you had {}",
            globals
                .names()
                .map(|s| format!("`{}`", s.as_str()))
                .join(", ")
        );

        *globals.0.variables.values().next().unwrap()
    }

    /// Move all the globals in this [`GlobalsBuilder`] into a new one. All variables will
    /// only be allocated once (ensuring things like function comparison works properly).
    pub fn populate(&'static self, x: impl FnOnce(&mut GlobalsBuilder), out: &mut GlobalsBuilder) {
        let globals = self.globals(x);
        for (name, value) in globals.0.variables.iter() {
            out.set(name.as_str(), *value)
        }
        out.docstring = globals.0.docstring.clone();
    }
}

impl MethodsStatic {
    /// Create a new [`MethodsStatic`].
    pub const fn new() -> Self {
        Self(OnceCell::new())
    }

    /// Populate the globals with a builder function. Always returns `Some`, but using this API
    /// to be a better fit for [`StarlarkValue.get_methods`](crate::values::StarlarkValue::get_methods).
    pub fn methods(&'static self, x: impl FnOnce(&mut MethodsBuilder)) -> Option<&'static Methods> {
        Some(self.0.get_or_init(|| MethodsBuilder::new().with(x).build()))
    }

    /// Move all the globals in this [`GlobalsBuilder`] into a new one. All variables will
    /// only be allocated once (ensuring things like function comparison works properly).
    pub fn populate(&'static self, x: impl FnOnce(&mut MethodsBuilder), out: &mut MethodsBuilder) {
        let methods = self.methods(x).unwrap();
        for (name, value) in methods.members.iter() {
            out.members.insert(name.as_str(), *value);
        }
        out.docstring = methods.docstring.clone();
    }
}

fn common_documentation<'a>(
    docstring: &Option<String>,
    members: impl IntoIterator<Item = (&'a str, FrozenValue)>,
) -> DocObject {
    let main_docs = docstring
        .as_ref()
        .and_then(|ds| DocString::from_docstring(DocStringKind::Rust, ds));
    let member_docs = members
        .into_iter()
        .map(|(name, val)| (name.to_owned(), DocMember::from_value(val.to_value())))
        .sorted_by(|(l, _), (r, _)| Ord::cmp(l, r))
        .collect();

    DocObject {
        docs: main_docs,
        members: member_docs,
    }
}

#[cfg(test)]
mod tests {
    use derive_more::Display;
    use starlark_derive::starlark_value;

    use super::*;
    use crate as starlark;
    use crate::any::ProvidesStaticType;
    use crate::assert::Assert;
    use crate::starlark_simple_value;
    use crate::values::NoSerialize;
    use crate::values::StarlarkValue;

    #[test]
    fn test_send_sync()
    where
        Globals: Send + Sync,
    {
    }

    #[test]
    fn test_set_attribute() {
        #[derive(Debug, Display, ProvidesStaticType, NoSerialize, Allocative)]
        #[display(fmt = "Magic")]
        struct Magic;
        starlark_simple_value!(Magic);

        #[starlark_value(type = "magic")]
        impl<'v> StarlarkValue<'v> for Magic {
            fn get_methods() -> Option<&'static Methods> {
                static RES: MethodsStatic = MethodsStatic::new();
                RES.methods(|x| {
                    x.set_attribute("my_type", "magic", None);
                    x.set_attribute("my_value", 42, None);
                })
            }
        }

        let mut a = Assert::new();
        a.globals_add(|x| x.set("magic", Magic));
        a.pass(
            r#"
assert_eq(magic.my_type, "magic")
assert_eq(magic.my_value, 42)"#,
        );
    }
}


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

//! Types representing Starlark modules ([`Module`] and [`FrozenModule`]) and global variables ([`Globals`]).
//!
//! Global functions and values are stored in [`Globals`], which are typically
//! built using [`GlobalsBuilder`].
//! User executions store their values in a [`Module`], which have to be converted to a
//! [`FrozenModule`] using [`freeze`](Module::freeze) before they can be `load()`'d as a dependency.

mod globals;
mod module_dump;
mod modules;
pub(crate) mod names;
pub(crate) mod slots;

pub use globals::*;
pub use modules::*;
use thiserror::Error;

#[derive(Debug, Error)]
enum EnvironmentError {
    /// Cannot import private symbol, i.e. underscore prefixed
    #[error("Cannot import private symbol `{0}`")]
    CannotImportPrivateSymbol(String),
    #[error("Module has no symbol `{0}`")]
    ModuleHasNoSymbol(String),
    #[error("Module has no symbol `{0}`, did you mean `{1}`?")]
    ModuleHasNoSymbolDidYouMean(String, String),
    #[error("Module symbol `{0}` is not exported")]
    ModuleSymbolIsNotExported(String),
}


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

//! The environment, called "Module" in [this spec](
//! https://github.com/google/skylark/blob/a0e5de7e63b47e716cca7226662a4c95d47bf873/doc/spec.md)
//! is the list of variable in the current scope. It can be frozen, after which
//! all values from this environment become immutable.

use std::cell::Cell;
use std::cell::RefCell;
use std::mem;
use std::time::Duration;
use std::time::Instant;

use allocative::Allocative;
use dupe::Dupe;
use itertools::Itertools;

use crate::cast::transmute;
use crate::collections::Hashed;
use crate::docs::DocMember;
use crate::docs::DocModule;
use crate::docs::DocString;
use crate::docs::DocStringKind;
use crate::environment::names::FrozenNames;
use crate::environment::names::MutableNames;
use crate::environment::slots::FrozenSlots;
use crate::environment::slots::ModuleSlotId;
use crate::environment::slots::MutableSlots;
use crate::environment::EnvironmentError;
use crate::environment::Globals;
use crate::errors::did_you_mean::did_you_mean;
use crate::eval::runtime::profile::heap::RetainedHeapProfileMode;
use crate::eval::ProfileData;
use crate::syntax::ast::Visibility;
use crate::values::layout::heap::heap_type::HeapKind;
use crate::values::layout::heap::profile::aggregated::AggregateHeapProfileInfo;
use crate::values::layout::heap::profile::aggregated::RetainedHeapProfile;
use crate::values::Freeze;
use crate::values::Freezer;
use crate::values::FrozenHeap;
use crate::values::FrozenHeapRef;
use crate::values::FrozenRef;
use crate::values::FrozenStringValue;
use crate::values::FrozenValue;
use crate::values::Heap;
use crate::values::OwnedFrozenValue;
use crate::values::Trace;
use crate::values::Tracer;
use crate::values::Value;

#[derive(Debug, thiserror::Error)]
enum ModuleError {
    #[error("Retained memory profiling is not enabled")]
    RetainedMemoryProfileNotEnabled,
}

/// The result of freezing a [`Module`], making it and its contained values immutable.
///
/// The values of this [`FrozenModule`] are stored on a frozen heap, a reference to which
/// can be obtained using [`frozen_heap`](FrozenModule::frozen_heap). Be careful not to use
/// these values after the [`FrozenModule`] has been released unless you obtain a reference
/// to the frozen heap.
#[derive(Debug, Clone, Dupe, Allocative)]
// We store the two elements separately since the FrozenHeapRef contains
// a copy of the FrozenModuleData inside it.
// Two Arc's should still be plenty cheap enough to qualify for `Dupe`.
pub struct FrozenModule {
    heap: FrozenHeapRef,
    module: FrozenRef<'static, FrozenModuleData>,
    extra_value: Option<FrozenValue>,
    /// Module evaluation duration:
    /// * evaluation of the top-level statements
    /// * optimizations during that evaluation
    /// * freezing and optimizations during freezing
    /// * does not include parsing time
    pub(crate) eval_duration: Duration,
}

#[derive(Debug, Allocative)]
pub(crate) struct FrozenModuleData {
    pub(crate) names: FrozenNames,
    pub(crate) slots: FrozenSlots,
    docstring: Option<String>,
    /// When heap profile enabled, this field stores retained memory info.
    heap_profile: Option<RetainedHeapProfile>,
}

/// A container for user values, used during execution.
///
/// A module contains both a [`FrozenHeap`] and [`Heap`] on which different values are allocated.
/// You can get references to these heaps with [`frozen_heap`](Module::frozen_heap) and
/// [`heap`](Module::heap). Be careful not to use these values after the [`Module`] has been
/// released unless you obtain a reference to the frozen heap.
#[derive(Debug)]
pub struct Module {
    heap: Heap,
    frozen_heap: FrozenHeap,
    names: MutableNames,
    // Should really be MutableSlots<'v>, where &'v self
    // Values are allocated from heap. Because of variance
    // you can inject the wrong values in, so make sure slots aren't
    // exported.
    slots: MutableSlots<'static>,
    docstring: RefCell<Option<String>>,
    /// Module evaluation duration:
    /// * evaluation of the top-level statements
    /// * optimizations during that evaluation
    /// * does not include freezing time
    /// * does not include parsing time
    eval_duration: Cell<Duration>,
    /// Field that can be used for any purpose you want.
    extra_value: Cell<Option<Value<'static>>>,
    /// When `Some`, heap profile is collected on freeze.
    heap_profile_on_freeze: Cell<Option<RetainedHeapProfileMode>>,
}

impl FrozenModule {
    /// Convert items in `globals` into a `FrozenModule`.
    /// This function can be used to implement starlark module
    /// using the `#[starlark_module]` attribute.
    ///
    /// This function does not return an error,
    /// but we prefer not to panic if there's some high level logic error.
    pub fn from_globals(globals: &Globals) -> anyhow::Result<FrozenModule> {
        let module = Module::new();

        module.frozen_heap.add_reference(globals.heap());

        for (name, value) in globals.iter() {
            module.set(name, value.to_value());
        }

        module.freeze()
    }

    fn get_any_visibility_option(&self, name: &str) -> Option<(OwnedFrozenValue, Visibility)> {
        self.module.names.get_name(name).and_then(|(slot, vis)|
        // This code is safe because we know the frozen module ref keeps the values alive
        self.module
            .slots
            .get_slot(slot)
            .map(|x| (unsafe { OwnedFrozenValue::new(self.heap.dupe(), x) }, vis)))
    }

    /// Get value, exported or private by name.
    // TODO(nga): separate private visibility into private (`_foo`) and imported (`load('foo')`).
    //   Users might want to access private variables, but not imported.
    #[doc(hidden)]
    pub fn get_any_visibility(&self, name: &str) -> anyhow::Result<(OwnedFrozenValue, Visibility)> {
        self.get_any_visibility_option(name).ok_or_else(|| {
            match did_you_mean(name, self.names().map(|s| s.as_str())) {
                Some(better) => EnvironmentError::ModuleHasNoSymbolDidYouMean(
                    name.to_owned(),
                    better.to_owned(),
                )
                .into(),
                None => EnvironmentError::ModuleHasNoSymbol(name.to_owned()).into(),
            }
        })
    }

    /// Get the value of the exported variable `name`.
    ///
    /// # Returns
    /// * `None` if symbol is not found
    /// * error if symbol is private
    pub fn get_option(&self, name: &str) -> anyhow::Result<Option<OwnedFrozenValue>> {
        match self.get_any_visibility_option(name) {
            None => Ok(None),
            Some((_, Visibility::Private)) => {
                Err(EnvironmentError::ModuleSymbolIsNotExported(name.to_owned()).into())
            }
            Some((value, Visibility::Public)) => Ok(Some(value)),
        }
    }

    /// Get the value of the exported variable `name`.
    /// Returns an error if the variable isn't defined in the module or it is private.
    pub fn get(&self, name: &str) -> anyhow::Result<OwnedFrozenValue> {
        self.get_any_visibility(name)
            .and_then(|(value, vis)| match vis {
                Visibility::Private => {
                    Err(EnvironmentError::ModuleSymbolIsNotExported(name.to_owned()).into())
                }
                Visibility::Public => Ok(value),
            })
    }

    /// Iterate through all the names defined in this module.
    /// Only includes symbols that are publicly exposed.
    pub fn names(&self) -> impl Iterator<Item = FrozenStringValue> + '_ {
        self.module.names()
    }

    /// Obtain the [`FrozenHeapRef`] which owns the storage of all values defined in this module.
    pub fn frozen_heap(&self) -> &FrozenHeapRef {
        &self.heap
    }

    /// Print out some approximation of the module definitions.
    pub fn describe(&self) -> String {
        self.module.describe()
    }

    pub(crate) fn all_items(&self) -> impl Iterator<Item = (FrozenStringValue, FrozenValue)> + '_ {
        self.module.all_items()
    }

    /// The documentation for the module, and all of its top level values
    ///
    /// Returns `(<module documentation>, { <symbol> : <that symbol's documentation> })`
    pub fn documentation(&self) -> DocModule {
        let members = self
            .all_items()
            .filter(|n| Module::default_visibility(n.0.as_str()) == Visibility::Public)
            .map(|(k, v)| (k.as_str().to_owned(), DocMember::from_value(v.to_value())))
            .collect();

        DocModule {
            docs: self.module.documentation(),
            members,
        }
    }

    /// Retained memory info, or error if not enabled.
    pub fn aggregated_heap_profile_info(&self) -> anyhow::Result<&AggregateHeapProfileInfo> {
        match &self.module.heap_profile {
            None => Err(ModuleError::RetainedMemoryProfileNotEnabled.into()),
            Some(p) => Ok(&p.info),
        }
    }

    /// Retained memory info, or error if not enabled.
    pub fn heap_profile(&self) -> anyhow::Result<ProfileData> {
        match &self.module.heap_profile {
            None => Err(ModuleError::RetainedMemoryProfileNotEnabled.into()),
            Some(p) => Ok(p.to_profile()),
        }
    }

    /// `extra_value` field from `Module`, frozen.
    pub fn extra_value(&self) -> Option<FrozenValue> {
        self.extra_value
    }

    /// `extra_value` field from `Module`, frozen.
    pub fn owned_extra_value(&self) -> Option<OwnedFrozenValue> {
        self.extra_value
            .map(|v| unsafe { OwnedFrozenValue::new(self.heap.dupe(), v) })
    }
}

impl FrozenModuleData {
    fn names(&self) -> impl Iterator<Item = FrozenStringValue> + '_ {
        self.names.symbols().map(|x| x.0)
    }

    fn describe(&self) -> String {
        self.items()
            .map(|(name, val)| val.to_value().describe(&name))
            .join("\n")
    }

    fn items(&self) -> impl Iterator<Item = (FrozenStringValue, FrozenValue)> + '_ {
        self.names
            .symbols()
            .filter_map(|(name, slot)| Some((name, self.slots.get_slot(slot)?)))
    }

    fn all_items(&self) -> impl Iterator<Item = (FrozenStringValue, FrozenValue)> + '_ {
        self.names
            .all_symbols()
            .filter_map(|(name, slot)| Some((name, self.slots.get_slot(slot)?)))
    }

    pub(crate) fn get_slot(&self, slot: ModuleSlotId) -> Option<FrozenValue> {
        self.slots.get_slot(slot)
    }

    /// Try and go back from a slot to a name.
    /// Inefficient - only use in error paths.
    pub(crate) fn get_slot_name(&self, slot: ModuleSlotId) -> Option<FrozenStringValue> {
        for (s, i) in self.names.symbols() {
            if i == slot {
                return Some(s);
            }
        }
        None
    }

    fn documentation(&self) -> Option<DocString> {
        self.docstring
            .as_ref()
            .and_then(|d| DocString::from_docstring(DocStringKind::Starlark, d))
    }
}

impl Default for Module {
    fn default() -> Self {
        Self::new()
    }
}

impl Module {
    /// Create a new module environment with no contents.
    pub fn new() -> Self {
        Self {
            heap: Heap::new(),
            frozen_heap: FrozenHeap::new(),
            names: MutableNames::new(),
            slots: MutableSlots::new(),
            docstring: RefCell::new(None),
            eval_duration: Cell::new(Duration::ZERO),
            extra_value: Cell::new(None),
            heap_profile_on_freeze: Cell::new(None),
        }
    }

    pub(crate) fn enable_heap_profile(&self, mode: RetainedHeapProfileMode) {
        self.heap_profile_on_freeze.set(Some(mode));
    }

    /// Get the heap on which values are allocated by this module.
    pub fn heap(&self) -> &Heap {
        &self.heap
    }

    /// Get the frozen heap on which frozen values are allocated by this module.
    pub fn frozen_heap(&self) -> &FrozenHeap {
        &self.frozen_heap
    }

    /// Iterate through all the names defined in this module.
    /// Only includes symbols that are publicly exposed.
    pub fn names(&self) -> impl Iterator<Item = FrozenStringValue> + '_ {
        self.names
            .all_names_and_visibilities()
            .into_iter()
            .filter_map(|(name, vis)| {
                if vis == Visibility::Public {
                    Some(name)
                } else {
                    None
                }
            })
    }

    /// Iterate through all the names defined in this module, including those that are private.
    pub fn names_and_visibilities(
        &self,
    ) -> impl Iterator<Item = (FrozenStringValue, Visibility)> + '_ {
        self.names.all_names_and_visibilities().into_iter()
    }

    pub(crate) fn mutable_names(&self) -> &MutableNames {
        &self.names
    }

    pub(crate) fn slots<'v>(&'v self) -> &'v MutableSlots<'v> {
        // Not true because of variance, but mostly true. Don't export further.
        unsafe { transmute!(&'v MutableSlots<'static>, &'v MutableSlots<'v>, &self.slots) }
    }

    /// Get value, exported or private by name.
    pub(crate) fn get_any_visibility<'v>(
        &'v self,
        name: Hashed<&str>,
    ) -> Option<(Value<'v>, Visibility)> {
        let (slot, vis) = self.names.get_name(name)?;
        let value = self.slots().get_slot(slot)?;
        Some((value, vis))
    }

    /// Get the value of the exported variable `name`.
    /// Returns [`None`] if the variable isn't defined in the module or it is private.
    pub fn get<'v>(&'v self, name: &str) -> Option<Value<'v>> {
        self.get_any_visibility(Hashed::new(name))
            .and_then(|(v, vis)| match vis {
                Visibility::Private => None,
                Visibility::Public => Some(v),
            })
    }

    /// Freeze the environment, all its value will become immutable afterwards.
    pub fn freeze(self) -> anyhow::Result<FrozenModule> {
        let Module {
            names,
            slots,
            frozen_heap,
            heap,
            docstring,
            eval_duration,
            extra_value,
            heap_profile_on_freeze,
        } = self;
        let start = Instant::now();
        // This is when we do the GC/freeze, using the module slots as roots
        // Note that we even freeze anonymous slots, since they are accessed by
        // slot-index in the code, and we don't walk into them, so don't know if
        // they are used.
        let freezer = Freezer::new(frozen_heap);
        let slots = slots.freeze(&freezer)?;
        let extra_value = extra_value.into_inner().freeze(&freezer)?;
        let stacks = if let Some(mode) = heap_profile_on_freeze.get() {
            // TODO(nga): retained heap profile does not store information about data
            //   allocated in frozen heap before freeze starts.
            let heap_profile = AggregateHeapProfileInfo::collect(&heap, Some(HeapKind::Frozen));
            Some(RetainedHeapProfile {
                info: heap_profile,
                mode,
            })
        } else {
            None
        };
        let rest = FrozenModuleData {
            names: names.freeze(),
            slots,
            docstring: docstring.into_inner(),
            heap_profile: stacks,
        };
        let frozen_module_ref = freezer.heap.alloc_any_display_from_debug(rest);
        for frozen_def in freezer.frozen_defs.borrow().as_slice() {
            frozen_def.post_freeze(frozen_module_ref, &heap, &freezer.heap);
        }
        // The values MUST be alive up until this point (as the above line uses them),
        // but can now be dropped
        mem::drop(heap);

        if let Some(stacks) = &frozen_module_ref.heap_profile {
            assert_eq!(stacks.info.unused_capacity.get(), 0, "sanity check");
            stacks
                .info
                .unused_capacity
                .set(freezer.heap.unused_capacity());
        }

        Ok(FrozenModule {
            heap: freezer.into_ref(),
            module: frozen_module_ref,
            extra_value,
            eval_duration: start.elapsed() + eval_duration.get(),
        })
    }

    /// Set the value of a variable in the environment.
    /// Modifying these variables while executing is ongoing can have
    /// surprising effects.
    pub fn set<'v>(&'v self, name: &str, value: Value<'v>) {
        let slot = self.names.add_name(self.frozen_heap.alloc_str_intern(name));
        let slots = self.slots();
        slots.ensure_slot(slot);
        slots.set_slot(slot, value);
    }

    /// Symbols starting with underscore are considered private.
    pub(crate) fn default_visibility(symbol: &str) -> Visibility {
        match symbol.starts_with('_') {
            true => Visibility::Private,
            false => Visibility::Public,
        }
    }

    /// Set the value of a variable in the environment. Set its visibliity to
    /// "private" to ensure that it is not re-exported
    pub(crate) fn set_private<'v>(&'v self, name: FrozenStringValue, value: Value<'v>) {
        let slot = self.names.add_name_visibility(name, Visibility::Private);
        let slots = self.slots();
        slots.ensure_slot(slot);
        slots.set_slot(slot, value);
    }

    /// Import symbols from a module, similar to what is done during `load()`.
    pub fn import_public_symbols(&self, module: &FrozenModule) {
        self.frozen_heap.add_reference(&module.heap);
        for (k, slot) in module.module.names.symbols() {
            if Self::default_visibility(&k) == Visibility::Public {
                if let Some(value) = module.module.slots.get_slot(slot) {
                    self.set_private(k, Value::new_frozen(value))
                }
            }
        }
    }

    pub(crate) fn load_symbol<'v>(
        &'v self,
        module: &FrozenModule,
        symbol: &str,
    ) -> anyhow::Result<Value<'v>> {
        if Self::default_visibility(symbol) != Visibility::Public {
            return Err(EnvironmentError::CannotImportPrivateSymbol(symbol.to_owned()).into());
        }
        match module.get_any_visibility(symbol)? {
            (v, Visibility::Public) => Ok(v.owned_value(self.frozen_heap())),
            (_, Visibility::Private) => {
                Err(EnvironmentError::ModuleSymbolIsNotExported(symbol.to_owned()).into())
            }
        }
    }

    pub(crate) fn set_docstring(&self, docstring: String) {
        self.docstring.replace(Some(docstring));
    }

    pub(crate) fn add_eval_duration(&self, duration: Duration) {
        self.eval_duration.set(self.eval_duration.get() + duration);
    }

    pub(crate) fn trace<'v>(&'v self, tracer: &Tracer<'v>) {
        self.slots().get_slots_mut().trace(tracer);

        let extra_value = self.extra_value();
        if let Some(mut extra_value) = extra_value {
            extra_value.trace(tracer);
            self.set_extra_value(extra_value);
        }
    }

    /// Field that can be used for any purpose you want.
    pub fn set_extra_value<'v>(&'v self, v: Value<'v>) {
        // Cast lifetime.
        let v = unsafe { transmute!(Value, Value, v) };
        self.extra_value.set(Some(v));
    }

    /// Field that can be used for any purpose you want.
    pub fn extra_value<'v>(&'v self) -> Option<Value<'v>> {
        // Cast lifetime.
        unsafe { transmute!(Option<Value>, Option<Value>, self.extra_value.get()) }
    }
}

#[test]
fn test_send_sync()
where
    FrozenModule: Send + Sync,
{
}

#[cfg(test)]
mod tests {
    use starlark_derive::starlark_module;

    use crate as starlark;
    use crate::environment::FrozenModule;
    use crate::environment::Globals;
    use crate::environment::GlobalsBuilder;
    use crate::environment::Module;
    use crate::eval::Evaluator;
    use crate::eval::ProfileMode;
    use crate::syntax::AstModule;
    use crate::syntax::Dialect;
    use crate::values::list::ListRef;

    #[test]
    fn test_gen_heap_summary_profile() {
        let module = Module::new();
        {
            let mut eval = Evaluator::new(&module);
            eval.enable_profile(&ProfileMode::HeapSummaryRetained)
                .unwrap();
            eval.eval_module(
                AstModule::parse(
                    "x.star",
                    r"
def f(x):
    return list([x])

x = f(1)
"
                    .to_owned(),
                    &Dialect::Extended,
                )
                .unwrap(),
                &Globals::standard(),
            )
            .unwrap();
        }
        let module = module.freeze().unwrap();
        let profile_info = module.aggregated_heap_profile_info().unwrap();
        let heap_summary = profile_info.gen_summary_csv();
        // Smoke test.
        assert!(profile_info.unused_capacity.get() > 0);
        assert!(heap_summary.contains("\"x.star.f\""), "{:?}", heap_summary);
    }

    #[test]
    fn test_frozen_module_from_globals() {
        #[starlark_module]
        fn some_globals(globals: &mut GlobalsBuilder) {
            fn foo() -> anyhow::Result<i32> {
                Ok(17)
            }

            const BAR: Vec<String> = Vec::new();
        }

        let mut globals = GlobalsBuilder::new();
        some_globals(&mut globals);
        let globals = globals.build();

        let module = FrozenModule::from_globals(&globals).unwrap();
        assert_eq!("function", module.get("foo").unwrap().value().get_type());
        assert_eq!(
            0,
            ListRef::from_value(module.get("BAR").unwrap().value())
                .unwrap()
                .len()
        );
    }
}


/*
 * 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::fmt::Write;

use crate::environment::FrozenModule;
use crate::eval::compiler::def::FrozenDef;
use crate::values::FrozenHeapRef;
use crate::values::FrozenValueTyped;

impl FrozenModule {
    /// Print a lot of module internals for debugging.
    pub fn dump_debug(&self) -> String {
        let mut w = String::new();

        writeln!(w, "Eval duration: {:.3}s", self.eval_duration.as_secs_f64()).unwrap();
        writeln!(w, "Heap stats:").unwrap();
        w.push_str(&self.frozen_heap().dump_debug());

        for (name, value) in self.all_items() {
            // TODO(nga): this prints public, private and imported symbols.
            //   We only care about public and private symbols, but no imported.
            writeln!(w).unwrap();
            writeln!(w, "{} = {}", name, value).unwrap();
            if let Some(def) = FrozenValueTyped::<FrozenDef>::new(value) {
                def.dump_debug()
                    .lines()
                    .for_each(|line| writeln!(w, "  {}", line).unwrap());
            }
        }
        w
    }
}

impl FrozenHeapRef {
    fn dump_debug(&self) -> String {
        let mut w = String::new();
        writeln!(w, "Allocated bytes: {}", self.allocated_bytes()).unwrap();
        writeln!(w, "Available bytes: {}", self.available_bytes()).unwrap();
        w
    }
}


/*
 * 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 allocative::Allocative;

use crate::collections::Hashed;
use crate::collections::SmallMap;
use crate::environment::slots::ModuleSlotId;
use crate::syntax::ast::Visibility;
use crate::values::FrozenStringValue;

/// MutableNames are how we allocate slots (index-based) to variables
/// (name-based). The slots field is the current active mapping of names to
/// index.
///
/// In a statement context there are things that define variables, e.g. x=...,
/// for x in ... Importantly, the expression x can refer to either a global x,
/// or a local x that hasn't yet been defined, based on the future
/// presence/absence of a statement defining x. Therefore, we first capture all
/// the definitions with collect_defines_lvalue, allocate them slots,
/// then replace variables with slot numbers when compiling.
///
/// Comprehensions are a bit different. Given [x for x in y] that defines x, but
/// in a way that shadows any existing x, and the definition immediately binds
/// x. We do that with add_scoped()/unscope(). On an add_scope, we allocate
/// fresh slots at the end, and bind them to the names in the comprehension.
/// On an unscope, we do the reverse, putting things back to how they were
/// before (apart from the total) number of slots required.
#[derive(Debug)]
pub(crate) struct MutableNames(RefCell<SmallMap<FrozenStringValue, (ModuleSlotId, Visibility)>>);

#[derive(Debug, Allocative)]
pub(crate) struct FrozenNames(SmallMap<FrozenStringValue, (ModuleSlotId, Visibility)>);

impl MutableNames {
    pub(crate) fn new() -> Self {
        Self(RefCell::new(SmallMap::new()))
    }

    pub(crate) fn slot_count(&self) -> u32 {
        self.0.borrow().len().try_into().unwrap()
    }

    /// Try and go back from a slot to a name.
    /// Inefficient - only use in error paths.
    pub(crate) fn get_slot(&self, slot: ModuleSlotId) -> Option<FrozenStringValue> {
        for (s, (i, _vis)) in &*self.0.borrow() {
            if *i == slot {
                return Some(*s);
            }
        }
        None
    }

    pub(crate) fn get_name(&self, name: Hashed<&str>) -> Option<(ModuleSlotId, Visibility)> {
        self.0.borrow().get_hashed(name).copied()
    }

    /// Add a name with explicit visibility to the module.
    pub(crate) fn add_name_visibility(
        &self,
        name: FrozenStringValue,
        vis: Visibility,
    ) -> ModuleSlotId {
        let mut x = self.0.borrow_mut();
        match x.get_mut_hashed(name.get_hashed().as_ref()) {
            Some((slot, stored_vis)) => {
                // Public visibility wins.
                if *stored_vis == Visibility::Private {
                    *stored_vis = vis;
                }
                *slot
            }
            None => {
                let slot = ModuleSlotId::new(x.len().try_into().unwrap());
                x.insert_hashed(name.get_hashed(), (slot, vis));
                slot
            }
        }
    }

    // Add an exported name, or if it's already there, return the existing name
    pub(crate) fn add_name(&self, name: FrozenStringValue) -> ModuleSlotId {
        self.add_name_visibility(name, Visibility::Public)
    }

    pub(crate) fn hide_name(&self, name: &str) {
        self.0.borrow_mut().remove(name);
    }

    pub(crate) fn all_names_and_slots(&self) -> Vec<(FrozenStringValue, ModuleSlotId)> {
        self.0
            .borrow()
            .iter()
            .map(|(name, (slot, _vis))| (*name, *slot))
            .collect()
    }

    pub(crate) fn all_names_and_visibilities(&self) -> Vec<(FrozenStringValue, Visibility)> {
        self.0
            .borrow()
            .iter()
            .map(|(name, (_slot, vis))| (*name, *vis))
            .collect()
    }

    pub(crate) fn freeze(self) -> FrozenNames {
        FrozenNames(self.0.into_inner())
    }
}

impl FrozenNames {
    pub(crate) fn get_name(&self, name: &str) -> Option<(ModuleSlotId, Visibility)> {
        self.0.get(name).copied()
    }

    /// Symbols including private.
    pub(crate) fn all_symbols(
        &self,
    ) -> impl Iterator<Item = (FrozenStringValue, ModuleSlotId)> + '_ {
        self.0.iter().map(|(name, (slot, _vis))| (*name, *slot))
    }

    /// Exported symbols.
    pub(crate) fn symbols(&self) -> impl Iterator<Item = (FrozenStringValue, ModuleSlotId)> + '_ {
        self.0.iter().filter_map(|(name, (slot, vis))| match vis {
            Visibility::Private => None,
            Visibility::Public => Some((*name, *slot)),
        })
    }
}


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

use std::cell::RefCell;
use std::cell::RefMut;

use allocative::Allocative;
use dupe::Dupe;

use crate::values::Freeze;
use crate::values::Freezer;
use crate::values::FrozenValue;
use crate::values::Value;

#[derive(Clone, Copy, Dupe, Debug, PartialEq, Eq, Allocative)]
pub(crate) struct ModuleSlotId(pub(crate) u32);

impl ModuleSlotId {
    pub fn new(index: u32) -> Self {
        Self(index)
    }
}

// Indexed slots of a module. May contain unassigned values as `None`.
#[derive(Debug)]
pub(crate) struct MutableSlots<'v>(RefCell<Vec<Option<Value<'v>>>>);

// Indexed slots of a module. May contain unassigned values as `None`.
#[derive(Debug, Allocative)]
pub(crate) struct FrozenSlots(Vec<Option<FrozenValue>>);

impl<'v> MutableSlots<'v> {
    pub fn new() -> Self {
        Self(RefCell::new(Vec::new()))
    }

    pub(crate) fn get_slots_mut(&self) -> RefMut<Vec<Option<Value<'v>>>> {
        self.0.borrow_mut()
    }

    pub fn get_slot(&self, slot: ModuleSlotId) -> Option<Value<'v>> {
        self.0.borrow()[slot.0 as usize]
    }

    pub fn set_slot(&self, slot: ModuleSlotId, value: Value<'v>) {
        self.0.borrow_mut()[slot.0 as usize] = Some(value);
    }

    pub fn ensure_slot(&self, slot: ModuleSlotId) {
        // To ensure that `slot` exists, we need at least `slot + 1` slots.
        self.ensure_slots(slot.0 + 1);
    }

    pub fn ensure_slots(&self, count: u32) {
        let mut slots = self.0.borrow_mut();
        if slots.len() >= count as usize {
            return;
        }
        let extra = count as usize - slots.len();
        slots.reserve(extra);
        for _ in 0..extra {
            slots.push(None);
        }
    }

    pub(crate) fn freeze(self, freezer: &Freezer) -> anyhow::Result<FrozenSlots> {
        let slots = self.0.into_inner().freeze(freezer)?;
        Ok(FrozenSlots(slots))
    }
}

impl FrozenSlots {
    pub fn get_slot(&self, slot: ModuleSlotId) -> Option<FrozenValue> {
        self.0[slot.0 as usize]
    }
}