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

//! Bytecode profiler.

use std::collections::HashMap;
use std::iter::Sum;
use std::mem;
use std::ops::AddAssign;

use dupe::Dupe;

use crate::eval::bc::opcode::BcOpcode;
use crate::eval::runtime::profile::csv::CsvWriter;
use crate::eval::runtime::profile::data::ProfileDataImpl;
use crate::eval::ProfileData;
use crate::eval::ProfileMode;

#[derive(Debug, thiserror::Error)]
enum BcProfileError {
    #[error("Can't call `write_bc_profile` unless you first call `enable_bc_profile`.")]
    BcProfilingNotEnabled,
}

#[derive(Default, Clone, Dupe, Copy, Debug)]
struct BcInstrStat {
    count: u64,
}

impl<'a> Sum<&'a BcInstrStat> for BcInstrStat {
    fn sum<I: Iterator<Item = &'a BcInstrStat>>(iter: I) -> Self {
        let mut sum = BcInstrStat::default();
        for BcInstrStat { count } in iter {
            sum.count += *count;
        }
        sum
    }
}

impl<'a> AddAssign<&'a BcInstrStat> for BcInstrStat {
    fn add_assign(&mut self, other: &'a BcInstrStat) {
        let BcInstrStat { count } = other;
        self.count += count;
    }
}

#[derive(Default, Clone, Copy, Dupe, Debug)]
struct BcInstrPairsStat {
    count: u64,
    // We are not measuring time here, because even time for single opcode
    // is not very accurate or helpful, and time for pairs is even less helpful.
}

impl<'a> AddAssign<&'a BcInstrPairsStat> for BcInstrPairsStat {
    fn add_assign(&mut self, other: &'a BcInstrPairsStat) {
        let BcInstrPairsStat { count } = other;
        self.count += count;
    }
}

#[derive(Clone, Debug)]
pub(crate) struct BcProfileData {
    by_instr: [BcInstrStat; BcOpcode::COUNT],
}

impl<'a> AddAssign<&'a BcProfileData> for BcProfileData {
    fn add_assign(&mut self, rhs: &'a BcProfileData) {
        for (lhs, rhs) in self.by_instr.iter_mut().zip(rhs.by_instr.iter()) {
            *lhs += rhs;
        }
    }
}

#[derive(Default, Clone, Debug)]
pub(crate) struct BcPairsProfileData {
    last: Option<BcOpcode>,
    by_instr: HashMap<[BcOpcode; 2], BcInstrPairsStat>,
}

impl<'a> AddAssign<&'a BcPairsProfileData> for BcPairsProfileData {
    fn add_assign(&mut self, rhs: &'a BcPairsProfileData) {
        self.last = None;
        for (pair, stat) in &rhs.by_instr {
            *self.by_instr.entry(*pair).or_default() += stat;
        }
    }
}

// Derive doesn't work here.
impl Default for BcProfileData {
    fn default() -> Self {
        BcProfileData {
            by_instr: [BcInstrStat::default(); BcOpcode::COUNT],
        }
    }
}

impl BcProfileData {
    fn before_instr(&mut self, opcode: BcOpcode) {
        self.by_instr[opcode as usize].count += 1;
    }

    pub(crate) fn gen_csv(&self) -> String {
        let mut by_instr: Vec<_> = self
            .by_instr
            .iter()
            .enumerate()
            .map(|(i, st)| (BcOpcode::by_number(i as u32).unwrap(), st))
            .collect();
        by_instr.sort_by_key(|(_opcode, st)| u64::MAX - st.count);
        let total: BcInstrStat = by_instr.iter().map(|(_opcode, st)| *st).sum();
        let mut csv = CsvWriter::new(["Opcode", "Count", "Count / Total"]);
        {
            csv.write_display("TOTAL");
            csv.write_value(total.count);
            csv.write_display(format!("{:.3}", 1.0));
            csv.finish_row();
        }
        for (opcode, instr_stats) in &by_instr {
            csv.write_debug(opcode);
            csv.write_value(instr_stats.count);
            csv.write_display(format!(
                "{:.3}",
                instr_stats.count as f64 / total.count as f64
            ));
            csv.finish_row();
        }
        csv.finish()
    }

    pub(crate) fn merge<'a>(iter: impl IntoIterator<Item = &'a BcProfileData>) -> BcProfileData {
        let mut sum = BcProfileData::default();
        for profile in iter {
            sum += profile;
        }
        sum
    }
}

impl BcPairsProfileData {
    fn before_instr(&mut self, opcode: BcOpcode) {
        if let Some(last_opcode) = self.last {
            self.by_instr
                .entry([last_opcode, opcode])
                .or_default()
                .count += 1;
        }
        self.last = Some(opcode);
    }

    pub(crate) fn gen_csv(&self) -> String {
        let mut by_instr: Vec<_> = self
            .by_instr
            .iter()
            .map(|(opcodes, stat)| (*opcodes, stat))
            .collect();
        by_instr.sort_by_key(|(opcodes, st)| (u64::MAX - st.count, *opcodes));
        let count_total = by_instr.iter().map(|(_, st)| st.count).sum::<u64>();
        let mut csv = CsvWriter::new(["Opcode[0]", "Opcode[1]", "Count", "Count / Total"]);
        for ([o0, o1], instr_stats) in &by_instr {
            csv.write_debug(o0);
            csv.write_debug(o1);
            csv.write_value(instr_stats.count);
            csv.write_display(format!(
                "{:.3}",
                instr_stats.count as f64 / count_total as f64
            ));
            csv.finish_row();
        }
        csv.finish()
    }

    pub(crate) fn merge<'a>(
        iter: impl IntoIterator<Item = &'a BcPairsProfileData>,
    ) -> BcPairsProfileData {
        let mut sum = BcPairsProfileData::default();
        for profile in iter {
            sum += profile;
        }
        sum
    }
}

enum BcProfileDataMode {
    Bc(Box<BcProfileData>),
    BcPairs(Box<BcPairsProfileData>),
    Disabled,
}

pub(crate) struct BcProfile {
    data: BcProfileDataMode,
}

impl BcProfile {
    pub(crate) fn new() -> BcProfile {
        BcProfile {
            data: BcProfileDataMode::Disabled,
        }
    }

    pub(crate) fn enable_1(&mut self) {
        self.data = BcProfileDataMode::Bc(Default::default());
    }

    pub(crate) fn enable_2(&mut self) {
        self.data = BcProfileDataMode::BcPairs(Default::default());
    }

    pub(crate) fn enabled(&self) -> bool {
        match self.data {
            BcProfileDataMode::Bc(..) => true,
            BcProfileDataMode::BcPairs(..) => true,
            BcProfileDataMode::Disabled => false,
        }
    }

    pub(crate) fn gen_bc_profile(&mut self) -> anyhow::Result<ProfileData> {
        match mem::replace(&mut self.data, BcProfileDataMode::Disabled) {
            BcProfileDataMode::Bc(bc) => Ok(ProfileData {
                profile_mode: ProfileMode::Bytecode,
                profile: ProfileDataImpl::Bc(bc),
            }),
            _ => Err(BcProfileError::BcProfilingNotEnabled.into()),
        }
    }

    pub(crate) fn gen_bc_pairs_profile(&mut self) -> anyhow::Result<ProfileData> {
        match mem::replace(&mut self.data, BcProfileDataMode::Disabled) {
            BcProfileDataMode::BcPairs(bc_pairs) => Ok(ProfileData {
                profile_mode: ProfileMode::BytecodePairs,
                profile: ProfileDataImpl::BcPairs(*bc_pairs),
            }),
            _ => Err(BcProfileError::BcProfilingNotEnabled.into()),
        }
    }

    /// Called from bytecode.
    pub(crate) fn before_instr(&mut self, opcode: BcOpcode) {
        match &mut self.data {
            BcProfileDataMode::Bc(data) => data.before_instr(opcode),
            BcProfileDataMode::BcPairs(data) => data.before_instr(opcode),
            BcProfileDataMode::Disabled => {
                unreachable!("this code is unreachable when bytecode profiling is not enabled")
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::environment::Globals;
    use crate::environment::Module;
    use crate::eval::bc::opcode::BcOpcode;
    use crate::eval::runtime::profile::bc::BcPairsProfileData;
    use crate::eval::runtime::profile::bc::BcProfileData;
    use crate::eval::Evaluator;
    use crate::eval::ProfileMode;
    use crate::syntax::AstModule;
    use crate::syntax::Dialect;

    #[test]
    fn test_smoke() {
        let module = Module::new();
        let globals = Globals::standard();
        let mut eval = Evaluator::new(&module);
        eval.enable_profile(&ProfileMode::Bytecode).unwrap();
        eval.eval_module(
            AstModule::parse("bc.star", "repr([1, 2])".to_owned(), &Dialect::Standard).unwrap(),
            &globals,
        )
        .unwrap();
        let csv = eval.gen_bc_profile().unwrap().gen().unwrap();
        assert!(
            csv.contains(&format!("\n\"{:?}\",1,", BcOpcode::CallFrozenNativePos)),
            "{:?}",
            csv
        );
    }

    #[test]
    fn test_smoke_2() {
        let module = Module::new();
        let globals = Globals::standard();
        let mut eval = Evaluator::new(&module);
        eval.enable_profile(&ProfileMode::BytecodePairs).unwrap();
        eval.eval_module(
            AstModule::parse("bc.star", "repr([1, 2])".to_owned(), &Dialect::Standard).unwrap(),
            &globals,
        )
        .unwrap();
        let csv = eval.gen_bc_pairs_profile().unwrap().gen().unwrap();
        assert!(
            csv.contains(&format!(
                "\n\"{:?}\",\"{:?}\",1",
                BcOpcode::ListOfConsts,
                BcOpcode::CallFrozenNativePos
            )),
            "{:?}",
            csv
        );
    }

    #[test]
    fn test_bc_profile_data_merge() {
        let bc = BcProfileData::default();
        // Smoke test.
        BcProfileData::merge([&bc, &bc, &bc]);
    }

    #[test]
    fn test_bc_pairs_profile_data_merge() {
        let bc = BcPairsProfileData::default();
        // Smoke test.
        BcPairsProfileData::merge([&bc, &bc, &bc]);
    }
}


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

//! Write CSV files.

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

use crate::eval::runtime::small_duration::SmallDuration;

/// Writer for CSV files.
pub(crate) struct CsvWriter {
    /// Column count in CSV file.
    column_count: usize,
    /// While writing a row, this is the current column index.
    current_column_index: usize,
    /// Write CSV there.
    buf: String,
}

fn quote_str_for_csv(s: &str) -> String {
    format!("\"{}\"", s.replace('\"', "\"\""))
}

impl CsvWriter {
    pub(crate) fn new<'c>(columns: impl IntoIterator<Item = &'c str>) -> CsvWriter {
        let mut buf = String::new();
        let mut column_count = 0;
        for (i, column) in columns.into_iter().enumerate() {
            if i != 0 {
                buf.push(',');
            }
            buf.push_str(column);
            column_count = i + 1;
        }
        buf.push('\n');
        CsvWriter {
            column_count,
            current_column_index: 0,
            buf,
        }
    }

    pub(crate) fn write_value(&mut self, value: impl CsvValue) {
        assert!(self.current_column_index < self.column_count);
        if self.current_column_index != 0 {
            self.buf.push(',');
        }
        self.buf.push_str(&value.format_for_csv());
        self.current_column_index += 1;
    }

    pub(crate) fn write_display(&mut self, value: impl Display) {
        struct Impl<V: Display>(V);

        impl<V: Display> CsvValue for Impl<V> {
            fn format_for_csv(&self) -> String {
                quote_str_for_csv(&self.0.to_string())
            }
        }

        self.write_value(Impl(value))
    }

    pub(crate) fn write_debug(&mut self, value: impl Debug) {
        struct Impl<V: Debug>(V);

        impl<V: Debug> CsvValue for Impl<V> {
            fn format_for_csv(&self) -> String {
                quote_str_for_csv(&format!("{:?}", &self.0))
            }
        }

        self.write_value(Impl(value))
    }

    pub(crate) fn finish_row(&mut self) {
        assert_eq!(self.current_column_index, self.column_count);
        self.current_column_index = 0;
        self.buf.push('\n');
    }

    pub(crate) fn finish(self) -> String {
        assert!(self.current_column_index == 0);
        self.buf
    }
}

pub(crate) trait CsvValue {
    fn format_for_csv(&self) -> String;
}

impl CsvValue for SmallDuration {
    fn format_for_csv(&self) -> String {
        format!("{:.3}", self.to_duration().as_secs_f64())
    }
}

impl<V: CsvValue + ?Sized> CsvValue for &'_ V {
    fn format_for_csv(&self) -> String {
        (*self).format_for_csv()
    }
}

impl CsvValue for str {
    fn format_for_csv(&self) -> String {
        quote_str_for_csv(self)
    }
}

impl CsvValue for usize {
    fn format_for_csv(&self) -> String {
        self.to_string()
    }
}

impl CsvValue for u64 {
    fn format_for_csv(&self) -> String {
        self.to_string()
    }
}

impl CsvValue for i32 {
    fn format_for_csv(&self) -> String {
        self.to_string()
    }
}

impl CsvValue for u128 {
    fn format_for_csv(&self) -> String {
        self.to_string()
    }
}

#[cfg(test)]
mod tests {
    use crate::eval::runtime::profile::csv::quote_str_for_csv;
    use crate::eval::runtime::profile::csv::CsvWriter;
    use crate::eval::runtime::small_duration::SmallDuration;

    #[test]
    fn test_csv_writer() {
        let mut csv = CsvWriter::new(["File", "Count", "Duration"]);
        csv.write_value("a.bzl");
        csv.write_value(10);
        csv.write_value(SmallDuration { nanos: 17_000_000 });
        csv.finish_row();
        csv.write_value("b.bzl");
        csv.write_value(20);
        csv.write_value(SmallDuration { nanos: 19_000_000 });
        csv.finish_row();
        assert_eq!(
            "\
File,Count,Duration
\"a.bzl\",10,0.017
\"b.bzl\",20,0.019
",
            csv.finish()
        )
    }

    #[test]
    fn test_quote_str_for_csv() {
        assert_eq!("\"a\"", quote_str_for_csv("a"));
        assert_eq!("\"a\"\"\"", quote_str_for_csv("a\""));
    }
}


/*
 * 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::fs;
use std::path::Path;

use anyhow::Context;
use dupe::Dupe;

use crate::eval::runtime::profile::bc::BcPairsProfileData;
use crate::eval::runtime::profile::bc::BcProfileData;
use crate::eval::runtime::profile::flamegraph::FlameGraphData;
use crate::eval::ProfileMode;
use crate::slice_vec_ext::SliceExt;
use crate::values::AggregateHeapProfileInfo;

#[derive(Debug, thiserror::Error)]
enum ProfileDataError {
    #[error("Profile data is not consistent with profile mode (internal error)")]
    ProfileDataNotConsistent,
    #[error("Empty profile list cannot be merged")]
    EmptyProfileList,
    #[error("Different profile modes in profile")]
    DifferentProfileModes,
    #[error("Merge of profile data for profile mode `{0}` is not implemented")]
    MergeNotImplemented(ProfileMode),
}

#[derive(Clone, Debug)]
pub(crate) enum ProfileDataImpl {
    Bc(Box<BcProfileData>),
    BcPairs(BcPairsProfileData),
    AggregateHeapProfileInfo(Box<AggregateHeapProfileInfo>),
    /// Flame graph data is in milliseconds.
    TimeFlameProfile(FlameGraphData),
    Other(String),
}

/// Collected profiling data.
#[derive(Clone, Debug)]
pub struct ProfileData {
    pub(crate) profile_mode: ProfileMode,
    /// Serialized to text (e.g. CSV or flamegraph).
    pub(crate) profile: ProfileDataImpl,
}

impl ProfileData {
    pub(crate) fn new(profile_mode: ProfileMode, profile: String) -> ProfileData {
        ProfileData {
            profile_mode,
            profile: ProfileDataImpl::Other(profile),
        }
    }

    /// Generate a string with profile data (e.g. CSV or flamegraph, depending on profile type).
    pub fn gen(&self) -> anyhow::Result<String> {
        match (&self.profile, &self.profile_mode) {
            (ProfileDataImpl::Other(profile), _) => Ok(profile.clone()),
            (ProfileDataImpl::Bc(bc), _) => Ok(bc.gen_csv()),
            (ProfileDataImpl::BcPairs(bc_pairs), _) => Ok(bc_pairs.gen_csv()),
            (
                ProfileDataImpl::AggregateHeapProfileInfo(profile),
                ProfileMode::HeapFlameRetained | ProfileMode::HeapFlameAllocated,
            ) => Ok(profile.gen_flame_graph()),
            (
                ProfileDataImpl::AggregateHeapProfileInfo(profile),
                ProfileMode::HeapSummaryRetained | ProfileMode::HeapSummaryAllocated,
            ) => Ok(profile.gen_summary_csv()),
            (ProfileDataImpl::AggregateHeapProfileInfo(_), _) => {
                Err(ProfileDataError::ProfileDataNotConsistent.into())
            }
            (ProfileDataImpl::TimeFlameProfile(data), ProfileMode::TimeFlame) => Ok(data.write()),
            (ProfileDataImpl::TimeFlameProfile(_), _) => {
                Err(ProfileDataError::ProfileDataNotConsistent.into())
            }
        }
    }

    /// Write to a file.
    pub fn write(&self, path: &Path) -> anyhow::Result<()> {
        fs::write(path, self.gen()?).with_context(|| {
            format!(
                "write profile `{}` data to `{}`",
                self.profile_mode,
                path.display()
            )
        })?;
        Ok(())
    }

    /// Merge profiles (aggregate).
    pub fn merge<'a>(
        profiles: impl IntoIterator<Item = &'a ProfileData>,
    ) -> anyhow::Result<ProfileData> {
        let profiles = Vec::from_iter(profiles);
        let profile_mode = match profiles.first() {
            None => return Err(ProfileDataError::EmptyProfileList.into()),
            Some(p) => p.profile_mode.dupe(),
        };
        for p in &profiles {
            if p.profile_mode != profile_mode {
                return Err(ProfileDataError::DifferentProfileModes.into());
            }
        }
        let profile = match &profile_mode {
            ProfileMode::Bytecode => {
                let profiles = profiles.try_map(|p| match &p.profile {
                    ProfileDataImpl::Bc(bc) => Ok(&**bc),
                    _ => Err(ProfileDataError::ProfileDataNotConsistent),
                })?;
                let profile = BcProfileData::merge(profiles);
                ProfileDataImpl::Bc(Box::new(profile))
            }
            ProfileMode::BytecodePairs => {
                let profiles = profiles.try_map(|p| match &p.profile {
                    ProfileDataImpl::BcPairs(bc_pairs) => Ok(bc_pairs),
                    _ => Err(ProfileDataError::ProfileDataNotConsistent),
                })?;
                let profile = BcPairsProfileData::merge(profiles);
                ProfileDataImpl::BcPairs(profile)
            }
            ProfileMode::HeapSummaryAllocated
            | ProfileMode::HeapSummaryRetained
            | ProfileMode::HeapFlameAllocated
            | ProfileMode::HeapFlameRetained => {
                let profiles = profiles.try_map(|p| match &p.profile {
                    ProfileDataImpl::AggregateHeapProfileInfo(profile) => Ok(&**profile),
                    _ => Err(ProfileDataError::ProfileDataNotConsistent),
                })?;
                let profile = AggregateHeapProfileInfo::merge(profiles);
                ProfileDataImpl::AggregateHeapProfileInfo(Box::new(profile))
            }
            ProfileMode::TimeFlame => {
                let profiles = profiles.try_map(|p| match &p.profile {
                    ProfileDataImpl::TimeFlameProfile(data) => Ok(data),
                    _ => Err(ProfileDataError::ProfileDataNotConsistent),
                })?;
                let profile = FlameGraphData::merge(profiles);
                ProfileDataImpl::TimeFlameProfile(profile)
            }
            profile_mode => {
                return Err(ProfileDataError::MergeNotImplemented(profile_mode.dupe()).into());
            }
        };
        Ok(ProfileData {
            profile_mode,
            profile,
        })
    }
}

#[cfg(test)]
mod tests {
    use dupe::Dupe;

    use crate::eval::runtime::profile::bc::BcPairsProfileData;
    use crate::eval::runtime::profile::data::ProfileDataImpl;
    use crate::eval::runtime::profile::flamegraph::FlameGraphData;
    use crate::eval::ProfileData;
    use crate::eval::ProfileMode;

    #[test]
    fn merge_bc() {
        let profile = ProfileData {
            profile_mode: ProfileMode::Bytecode,
            profile: ProfileDataImpl::Bc(Box::default()),
        };
        // Smoke.
        ProfileData::merge([&profile, &profile]).unwrap();
    }

    #[test]
    fn merge_bc_pairs() {
        let profile = ProfileData {
            profile_mode: ProfileMode::BytecodePairs,
            profile: ProfileDataImpl::BcPairs(BcPairsProfileData::default()),
        };
        // Smoke.
        ProfileData::merge([&profile, &profile]).unwrap();
    }

    #[test]
    fn merge_aggregated_heap_profile() {
        for profile_mode in [
            ProfileMode::HeapFlameRetained,
            ProfileMode::HeapFlameAllocated,
            ProfileMode::HeapSummaryRetained,
            ProfileMode::HeapSummaryAllocated,
        ] {
            let profile = ProfileData {
                profile_mode: profile_mode.dupe(),
                profile: ProfileDataImpl::AggregateHeapProfileInfo(Box::default()),
            };
            // Smoke.
            ProfileData::merge([&profile, &profile]).unwrap();
        }
    }

    #[test]
    fn merge_time_flame() {
        let profile = ProfileData {
            profile_mode: ProfileMode::TimeFlame,
            profile: ProfileDataImpl::TimeFlameProfile(FlameGraphData::default()),
        };
        // Smoke.
        ProfileData::merge([&profile, &profile]).unwrap();
    }
}


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

//! Utility to write files in formats understood by `flamegraph.pl`.

use std::fmt::Write;

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

use crate::values::layout::heap::profile::arc_str::ArcStr;

/// Node in flamegraph tree.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) struct FlameGraphNode {
    children: SmallMap<ArcStr, FlameGraphNode>,
    value: Option<u64>,
}

/// Profiling data as flame tree.
///
/// Can be written to `flamegraph.pl` format.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) struct FlameGraphData {
    root: FlameGraphNode,
}

impl FlameGraphNode {
    fn write<'a>(&'a self, writer: &mut FlameGraphWriter, stack: &mut Vec<&'a str>) {
        if let Some(value) = self.value {
            writer.write(stack.iter().copied(), value);
        }
        for (k, v) in self.children.iter() {
            stack.push(k);
            v.write(writer, stack);
            stack.pop().unwrap();
        }
    }

    /// Add value to the node.
    pub(crate) fn add(&mut self, value: u64) {
        match &mut self.value {
            None => self.value = Some(value),
            Some(v) => *v += value,
        }
    }

    pub(crate) fn merge(&mut self, other: &FlameGraphNode) {
        if let Some(value) = other.value {
            self.add(value);
        }

        for (k, v) in &other.children {
            self.child(k.dupe()).merge(v);
        }
    }

    /// Get or create a child node.
    pub(crate) fn child(&mut self, name: ArcStr) -> &mut FlameGraphNode {
        self.children.entry(name).or_default()
    }
}

impl FlameGraphData {
    pub(crate) fn write(&self) -> String {
        let mut writer = FlameGraphWriter::new();
        let mut stack = Vec::new();
        self.root.write(&mut writer, &mut stack);
        assert!(stack.is_empty());
        writer.finish()
    }

    pub(crate) fn root(&mut self) -> &mut FlameGraphNode {
        &mut self.root
    }

    pub(crate) fn merge<'a>(
        graphs: impl IntoIterator<Item = &'a FlameGraphData>,
    ) -> FlameGraphData {
        let mut result = FlameGraphData::default();
        for graph in graphs {
            result.root.merge(&graph.root);
        }
        result
    }
}

pub(crate) struct FlameGraphWriter {
    buf: String,
}

impl FlameGraphWriter {
    pub(crate) fn new() -> FlameGraphWriter {
        FlameGraphWriter { buf: String::new() }
    }

    pub(crate) fn write<'s>(&mut self, key: impl IntoIterator<Item = &'s str>, value: u64) {
        let key = key.into_iter().collect::<Vec<_>>();
        if key.is_empty() {
            writeln!(self.buf, "(unknown) {}", value).unwrap();
        } else {
            writeln!(self.buf, "{} {}", key.join(";"), value).unwrap();
        }
    }

    pub(crate) fn finish(self) -> String {
        self.buf
    }
}

#[cfg(test)]
mod tests {
    use crate::eval::runtime::profile::flamegraph::FlameGraphData;
    use crate::eval::runtime::profile::flamegraph::FlameGraphWriter;

    #[test]
    fn test_flamegraph_writer() {
        let mut writer = FlameGraphWriter::new();
        writer.write(["aa", "bb"], 20);
        assert_eq!("aa;bb 20\n", writer.finish());
    }

    #[test]
    fn test_flamegraph_data() {
        let mut data = FlameGraphData::default();
        data.root().child("a".into()).add(10);
        data.root().child("a".into()).child("b".into()).add(20);
        data.root().child("a".into()).add(30);
        let data = data.write();
        assert_eq!("a 40\na;b 20\n", data);
    }

    #[test]
    fn test_merge() {
        let mut a = FlameGraphData::default();
        a.root().add(10);
        a.root().child("a".into()).add(100);
        a.root().child("b".into()).child("c".into()).add(1000);
        let mut b = FlameGraphData::default();
        b.root().add(20);
        b.root().child("a".into()).add(200);

        let c = FlameGraphData::merge([&a, &b]);

        let mut expected = FlameGraphData::default();
        expected.root().add(30);
        expected.root().child("a".into()).add(300);
        expected
            .root()
            .child("b".into())
            .child("c".into())
            .add(1000);

        assert_eq!(expected, c);
    }
}


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

use std::fmt::Debug;

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

use crate::eval::runtime::profile::data::ProfileData;
use crate::eval::runtime::profile::data::ProfileDataImpl;
use crate::eval::ProfileMode;
use crate::values::layout::heap::profile::aggregated::AggregateHeapProfileInfo;
use crate::values::Heap;
use crate::values::Value;

#[derive(Copy, Clone, Dupe, Debug, Allocative)]
pub(crate) enum RetainedHeapProfileMode {
    Flame,
    Summary,
}

impl RetainedHeapProfileMode {
    pub(crate) fn to_profile_mode(self) -> ProfileMode {
        match self {
            RetainedHeapProfileMode::Flame => ProfileMode::HeapFlameRetained,
            RetainedHeapProfileMode::Summary => ProfileMode::HeapSummaryRetained,
        }
    }
}

#[derive(Debug, thiserror::Error)]
enum HeapProfileError {
    #[error("heap profile not enabled")]
    NotEnabled,
}

#[derive(Copy, Clone, Dupe, Debug)]
pub(crate) enum HeapProfileFormat {
    Summary,
    FlameGraph,
}

pub(crate) struct HeapProfile {
    enabled: bool,
}

impl HeapProfile {
    pub(crate) fn new() -> Self {
        Self { enabled: false }
    }

    pub(crate) fn enable(&mut self) {
        self.enabled = true;
    }

    #[cold]
    #[inline(never)]
    pub(crate) fn record_call_enter<'v>(&self, function: Value<'v>, heap: &'v Heap) {
        if self.enabled {
            heap.record_call_enter(function);
        }
    }

    #[cold]
    #[inline(never)]
    pub(crate) fn record_call_exit<'v>(&self, heap: &'v Heap) {
        if self.enabled {
            heap.record_call_exit();
        }
    }

    // We could expose profile on the Heap, but it's an implementation detail that it works here.
    pub(crate) fn gen(
        &self,
        heap: &Heap,
        format: HeapProfileFormat,
    ) -> anyhow::Result<ProfileData> {
        if !self.enabled {
            return Err(HeapProfileError::NotEnabled.into());
        }
        Ok(Self::gen_enabled(heap, format))
    }

    pub(crate) fn gen_enabled(heap: &Heap, format: HeapProfileFormat) -> ProfileData {
        match format {
            HeapProfileFormat::Summary => Self::write_summarized_heap_profile(heap),
            HeapProfileFormat::FlameGraph => Self::write_flame_heap_profile(heap),
        }
    }

    fn write_flame_heap_profile(heap: &Heap) -> ProfileData {
        let stacks = AggregateHeapProfileInfo::collect(heap, None);
        ProfileData {
            profile_mode: ProfileMode::HeapFlameAllocated,
            profile: ProfileDataImpl::AggregateHeapProfileInfo(Box::new(stacks)),
        }
    }

    fn write_summarized_heap_profile(heap: &Heap) -> ProfileData {
        let stacks = AggregateHeapProfileInfo::collect(heap, None);
        ProfileData {
            profile_mode: ProfileMode::HeapSummaryAllocated,
            profile: ProfileDataImpl::AggregateHeapProfileInfo(Box::new(stacks)),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::environment::Globals;
    use crate::environment::Module;
    use crate::eval::Evaluator;
    use crate::eval::ProfileMode;
    use crate::syntax::AstModule;
    use crate::syntax::Dialect;
    use crate::values::Value;

    #[test]
    fn test_profiling() -> anyhow::Result<()> {
        // We don't test that the profile looks any particular way, but we do test it doesn't crash
        let ast = AstModule::parse(
            "foo.bzl",
            r#"
def f(x):
    return (x * 5) + 3
y = 8 * 9 + 2
f
"#
            .to_owned(),
            &Dialect::Extended,
        )?;
        let globals = Globals::standard();
        let module = Module::new();
        let module2 = Module::new();
        let module3 = Module::new();

        let mut eval = Evaluator::new(&module);
        eval.enable_profile(&ProfileMode::HeapSummaryAllocated)
            .unwrap();
        let f = eval.eval_module(ast, &globals)?;

        // first check module profiling works
        HeapProfile::write_summarized_heap_profile(module.heap());
        HeapProfile::write_flame_heap_profile(module.heap());

        // second check function profiling works
        let mut eval = Evaluator::new(&module2);
        eval.enable_profile(&ProfileMode::HeapSummaryAllocated)
            .unwrap();
        eval.eval_function(f, &[Value::testing_new_int(100)], &[])?;

        HeapProfile::write_summarized_heap_profile(module2.heap());
        HeapProfile::write_flame_heap_profile(module2.heap());

        // finally, check a user can add values into the heap before/after
        let mut eval = Evaluator::new(&module3);
        module3.heap().alloc("Thing that goes before");
        eval.enable_profile(&ProfileMode::HeapSummaryAllocated)
            .unwrap();
        eval.eval_function(f, &[Value::testing_new_int(100)], &[])?;

        module3.heap().alloc("Thing that goes after");
        HeapProfile::write_summarized_heap_profile(module3.heap());
        HeapProfile::write_flame_heap_profile(module3.heap());

        Ok(())
    }
}


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

use std::fmt::Display;
use std::str::FromStr;

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

pub(crate) mod bc;
pub(crate) mod csv;
pub(crate) mod data;
pub(crate) mod flamegraph;
pub(crate) mod heap;
pub(crate) mod or_instrumentation;
pub(crate) mod stmt;
pub(crate) mod time_flame;
pub(crate) mod typecheck;

/// How to profile starlark code.
#[derive(Debug, PartialEq, Eq, Hash, Clone, Dupe, Allocative)]
#[non_exhaustive]
pub enum ProfileMode {
    /// The heap profile mode provides information about the time spent in each function and allocations
    /// performed by each function. Enabling this mode the side effect of disabling garbage-collection.
    /// This profiling mode is the recommended one.
    HeapSummaryAllocated,
    /// Like heap summary, but information about retained memory after module is frozen.
    HeapSummaryRetained,
    /// Like heap profile, but writes output comparible with
    /// [flamegraph.pl](https://github.com/brendangregg/FlameGraph/blob/master/flamegraph.pl).
    HeapFlameAllocated,
    /// Like heap flame, but information about retained memory after module is frozen.
    HeapFlameRetained,
    /// The statement profile mode provides information about time spent in each statement.
    Statement,
    /// Code coverage.
    Coverage,
    /// The bytecode profile mode provides information about bytecode instructions.
    Bytecode,
    /// The bytecode profile mode provides information about bytecode instruction pairs.
    BytecodePairs,
    /// Provide output compatible with
    /// [flamegraph.pl](https://github.com/brendangregg/FlameGraph/blob/master/flamegraph.pl).
    TimeFlame,
    /// Profile runtime typechecking.
    Typecheck,
}

impl Display for ProfileMode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.name())
    }
}

impl ProfileMode {
    fn name(&self) -> &str {
        match self {
            ProfileMode::HeapSummaryAllocated => "heap-summary-allocated",
            ProfileMode::HeapSummaryRetained => "heap-summary-retained",
            ProfileMode::HeapFlameAllocated => "heap-flame-allocated",
            ProfileMode::HeapFlameRetained => "heap-flame-retained",
            ProfileMode::Statement => "statement",
            ProfileMode::Coverage => "coverage",
            ProfileMode::Bytecode => "bytecode",
            ProfileMode::BytecodePairs => "bytecode-pairs",
            ProfileMode::TimeFlame => "time-flame",
            ProfileMode::Typecheck => "typecheck",
        }
    }
}

impl FromStr for ProfileMode {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        for mode in [
            ProfileMode::HeapSummaryAllocated,
            ProfileMode::HeapSummaryRetained,
            ProfileMode::HeapFlameAllocated,
            ProfileMode::HeapFlameRetained,
            ProfileMode::Statement,
            ProfileMode::Coverage,
            ProfileMode::Bytecode,
            ProfileMode::BytecodePairs,
            ProfileMode::TimeFlame,
            ProfileMode::Typecheck,
        ] {
            if s == mode.name() {
                return Ok(mode);
            }
        }
        Err(anyhow::anyhow!("Invalid ProfileMode: `{}`", s))
    }
}


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

use dupe::Dupe;

use crate::eval::ProfileMode;

#[derive(Debug, Default, Clone, Dupe, Eq, PartialEq)]
pub(crate) enum ProfileOrInstrumentationMode {
    #[default]
    None,
    Profile(ProfileMode),
    Collected,
}


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

use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::collections::HashSet;
use std::iter;
use std::time::Instant;

use dupe::Dupe;

use crate::codemap::CodeMap;
use crate::codemap::CodeMapId;
use crate::codemap::FileSpan;
use crate::codemap::FileSpanRef;
use crate::codemap::ResolvedFileSpan;
use crate::codemap::Span;
use crate::eval::runtime::profile::csv::CsvWriter;
use crate::eval::runtime::profile::data::ProfileData;
use crate::eval::runtime::small_duration::SmallDuration;
use crate::eval::ProfileMode;

#[derive(Debug, thiserror::Error)]
enum StmtProfileError {
    #[error("Statement profiling is not enabled")]
    NotEnabled,
}

// When line profiling is not enabled, we want this to be small and cheap
pub(crate) struct StmtProfile(Option<Box<StmtProfileData>>);

// So we don't need a special case for the first time around,
// we have a special FileId of empty that we ignore when printing
#[derive(Clone)]
struct StmtProfileData {
    files: HashMap<CodeMapId, CodeMap>,
    stmts: HashMap<(CodeMapId, Span), (usize, SmallDuration)>,
    next_file: CodeMapId,
    last_span: (CodeMapId, Span),
    last_start: Instant,
}

impl StmtProfileData {
    fn new() -> Self {
        StmtProfileData {
            files: HashMap::new(),
            stmts: HashMap::new(),
            next_file: CodeMapId::EMPTY,
            last_span: (CodeMapId::EMPTY, Span::default()),
            last_start: Instant::now(),
        }
    }

    // Add the data from last_span into the entries
    fn add_last(&mut self, now: Instant) {
        let time = now - self.last_start;
        match self.stmts.entry(self.last_span) {
            Entry::Occupied(mut x) => {
                let v = x.get_mut();
                v.0 += 1;
                v.1 += SmallDuration::from_duration(time);
            }
            Entry::Vacant(x) => {
                x.insert((1, SmallDuration::from_duration(time)));
            }
        }
    }

    fn before_stmt(&mut self, span: Span, codemap: &CodeMap) {
        let now = Instant::now();
        self.add_last(now);
        if self.last_span.0 != codemap.id() {
            self.add_codemap(codemap);
        }
        self.last_span = (self.next_file, span);
        self.last_start = now;
    }

    fn add_codemap(&mut self, codemap: &CodeMap) {
        let id = codemap.id();
        self.next_file = id;
        match self.files.entry(id) {
            Entry::Occupied(_) => {
                // Nothing to do, we have already got an owned version of this CodeMap
            }
            Entry::Vacant(x) => {
                x.insert(codemap.dupe());
            }
        }
    }

    fn write_to_string(&self, now: Instant) -> String {
        // The statement that was running last won't have been properly updated.
        // However, at this point, we have probably run some post-execution code,
        // so it probably wouldn't have a "fair" timing anyway.
        // We do our best though, and give it a time of now.
        // Clone first, since we don't want to impact the real timing with our odd
        // final execution finish.
        let mut data = self.clone();
        data.add_last(now);

        struct Item {
            span: FileSpan,
            time: SmallDuration,
            count: usize,
        }
        // There should be one EMPTY span entry
        let mut items = Vec::with_capacity(data.stmts.len() - 1);
        let mut total_time = SmallDuration::default();
        let mut total_count = 0;
        for ((file, span), (count, time)) in data.stmts {
            // EMPTY represents the first time special-case
            if file != CodeMapId::EMPTY {
                let span = data.files[&file].file_span(span);
                total_time += time;
                total_count += count;
                items.push(Item { span, time, count })
            }
        }
        items.sort_by_key(|x| -(x.time.nanos as i128));

        let mut csv = CsvWriter::new(["File", "Span", "Duration(s)", "Count"]);
        csv.write_value("TOTAL");
        csv.write_value("");
        csv.write_value(total_time);
        csv.write_value(total_count);
        csv.finish_row();

        for x in items {
            csv.write_value(x.span.file.filename());
            csv.write_display(x.span.file.resolve_span(x.span.span));
            csv.write_value(x.time);
            csv.write_value(x.count);
            csv.finish_row();
        }

        csv.finish()
    }

    fn coverage(&self) -> HashSet<ResolvedFileSpan> {
        self.stmts
            .keys()
            .filter(|(file, _)| *file != CodeMapId::EMPTY)
            .chain(iter::once(&self.last_span))
            .map(|(code_map_id, span)| {
                self.files
                    .get(code_map_id)
                    .unwrap()
                    .file_span(*span)
                    .resolve()
            })
            .collect()
    }
}

impl StmtProfile {
    pub(crate) fn new() -> Self {
        Self(None)
    }

    pub(crate) fn enable(&mut self) {
        self.0 = Some(Box::new(StmtProfileData::new()))
    }

    pub(crate) fn before_stmt(&mut self, span: FileSpanRef) {
        if let Some(data) = &mut self.0 {
            data.before_stmt(span.span, span.file)
        }
    }

    // None = not applicable because not enabled
    pub(crate) fn gen(&self) -> anyhow::Result<ProfileData> {
        let now = Instant::now();
        match &self.0 {
            Some(data) => Ok(ProfileData::new(
                ProfileMode::Statement,
                data.write_to_string(now),
            )),
            None => Err(StmtProfileError::NotEnabled.into()),
        }
    }

    pub(crate) fn coverage(&self) -> anyhow::Result<HashSet<ResolvedFileSpan>> {
        Ok(self
            .0
            .as_ref()
            .ok_or(StmtProfileError::NotEnabled)?
            .coverage())
    }
}

#[cfg(test)]
mod tests {

    use crate::assert::test_functions;
    use crate::environment::GlobalsBuilder;
    use crate::environment::Module;
    use crate::eval::Evaluator;
    use crate::eval::ProfileMode;
    use crate::syntax::AstModule;
    use crate::syntax::Dialect;

    #[test]
    fn test_coverage() {
        let module = Module::new();
        let mut eval = Evaluator::new(&module);

        let module = AstModule::parse(
            "cov.star",
            r#"
def xx(x):
    return noop(x)

xx(*[1])
xx(*[2])
"#
            .to_owned(),
            &Dialect::Extended,
        )
        .unwrap();
        eval.enable_profile(&ProfileMode::Coverage).unwrap();
        let mut globals = GlobalsBuilder::standard();
        test_functions(&mut globals);
        eval.eval_module(module, &globals.build()).unwrap();

        let mut coverage: Vec<String> = eval
            .coverage()
            .unwrap()
            .into_iter()
            .map(|s| s.to_string())
            .collect();
        coverage.sort();
        assert_eq!(
            [
                "cov.star:2:1-5:1",
                "cov.star:3:5-19",
                "cov.star:5:1-9",
                "cov.star:6:1-9"
            ]
            .as_slice(),
            coverage
        );
    }
}


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

use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::slice;
use std::time::Instant;

use dupe::Dupe;
use starlark_map::StarlarkHasherBuilder;

use crate as starlark;
use crate::eval::runtime::profile::data::ProfileData;
use crate::eval::runtime::profile::data::ProfileDataImpl;
use crate::eval::runtime::profile::flamegraph::FlameGraphData;
use crate::eval::runtime::profile::flamegraph::FlameGraphNode;
use crate::eval::runtime::small_duration::SmallDuration;
use crate::eval::ProfileMode;
use crate::slice_vec_ext::SliceExt;
use crate::values::layout::heap::profile::arc_str::ArcStr;
use crate::values::layout::pointer::RawPointer;
use crate::values::FrozenValue;
use crate::values::Trace;
use crate::values::Tracer;
use crate::values::Value;

#[derive(Debug, thiserror::Error)]
enum FlameProfileError {
    #[error("Flame profile not enabled")]
    NotEnabled,
}

#[derive(Hash, PartialEq, Eq, Clone, Copy, Dupe)]
struct MutableValueId(usize);
#[derive(Hash, PartialEq, Eq, Clone, Copy, Dupe)]
struct FrozenValueId(usize);

/// Index into FlameData.values
#[derive(Hash, PartialEq, Eq, Clone, Copy, Dupe)]
enum ValueId {
    // This struct is two words, can be made one.
    Mutable(MutableValueId),
    Frozen(FrozenValueId),
}

impl ValueId {
    fn lookup<'a, T>(self, mutable: &'a [T], frozen: &'a [T]) -> &'a T {
        match self {
            ValueId::Mutable(x) => &mutable[x.0],
            ValueId::Frozen(x) => &frozen[x.0],
        }
    }
}

/// Bimap between `Value` and `ValueId`.
/// In order to optimise GC (which otherwise quickly becomes O(n^2)) we have to
/// dedupe the values, so store them in `values`, with a fast map to get them in `map`.
/// Whenever we GC, regenerate map.
#[derive(Default)]
struct ValueIndex<'v> {
    /// Map from `MutableValueId` to `Value`.
    mutable_values: Vec<Value<'v>>,
    /// Map from `FrozenValueId` to `Value`.
    frozen_values: Vec<FrozenValue>,
    /// Map from `Value` to `MutableValueId`.
    mutable_map: HashMap<RawPointer, MutableValueId, StarlarkHasherBuilder>,
    /// Map from `Value` to `FrozenValueId`.
    frozen_map: HashMap<RawPointer, FrozenValueId, StarlarkHasherBuilder>,
}

unsafe impl<'v> Trace<'v> for ValueIndex<'v> {
    fn trace(&mut self, tracer: &Tracer<'v>) {
        // We only need to trace mutable values.
        self.mutable_values.trace(tracer);
        // Have to rebuild the map, as its keyed by ValuePtr which changes on GC
        self.mutable_map.clear();
        for (i, x) in self.mutable_values.iter().enumerate() {
            self.mutable_map.insert(x.ptr_value(), MutableValueId(i));
        }
    }
}

impl<'v> ValueIndex<'v> {
    /// Map `Value` to `ValueId`.
    fn index(&mut self, value: Value<'v>) -> ValueId {
        match value.unpack_frozen() {
            Some(frozen) => match self.frozen_map.entry(frozen.ptr_value()) {
                Entry::Occupied(e) => ValueId::Frozen(*e.get()),
                Entry::Vacant(e) => {
                    let res = FrozenValueId(self.frozen_values.len());
                    self.frozen_values.push(frozen);
                    e.insert(res);
                    ValueId::Frozen(res)
                }
            },
            None => match self.mutable_map.entry(value.ptr_value()) {
                Entry::Occupied(e) => ValueId::Mutable(*e.get()),
                Entry::Vacant(e) => {
                    let res = MutableValueId(self.mutable_values.len());
                    self.mutable_values.push(value);
                    e.insert(res);
                    ValueId::Mutable(res)
                }
            },
        }
    }
}

enum Frame {
    /// Entry recorded when we enter a function.
    Push(ValueId),
    /// Entry recorded when we exit a function.
    Pop,
}

#[derive(Trace)]
pub(crate) struct TimeFlameProfile<'v>(
    /// `Some` means enabled.
    Option<Box<FlameData<'v>>>,
);

#[derive(Default, Trace)]
struct FlameData<'v> {
    /// All events in the profile, i.e. function entry or exit with timestamp.
    frames: Vec<(Frame, Instant)>,
    index: ValueIndex<'v>,
}

struct Stacks<'a> {
    name: &'a str,
    time: SmallDuration,
    children: HashMap<ValueId, Stacks<'a>, StarlarkHasherBuilder>,
}

impl<'a> Stacks<'a> {
    fn blank(name: &'a str) -> Self {
        Stacks {
            name,
            time: SmallDuration::default(),
            children: HashMap::with_hasher(StarlarkHasherBuilder),
        }
    }

    fn new(
        mutable_names: &'a [String],
        frozen_names: &'a [String],
        frames: &[(Frame, Instant)],
    ) -> Self {
        let mut res = Stacks::blank("root");
        let Some(mut last_time) = frames.first().map(|x| x.1) else {
            return res;
        };
        res.add(
            mutable_names,
            frozen_names,
            &mut frames.iter(),
            &mut last_time,
        );
        res
    }

    fn add(
        &mut self,
        mutable_names: &'a [String],
        frozen_names: &'a [String],
        frames: &mut slice::Iter<(Frame, Instant)>,
        last_time: &mut Instant,
    ) {
        while let Some((frame, time)) = frames.next() {
            self.time += time.duration_since(*last_time);
            *last_time = *time;
            match frame {
                Frame::Pop => return,
                Frame::Push(i) => match self.children.entry(*i) {
                    Entry::Occupied(mut e) => {
                        e.get_mut()
                            .add(mutable_names, frozen_names, frames, last_time)
                    }
                    Entry::Vacant(e) => e
                        .insert(Stacks::blank(
                            i.lookup(mutable_names, frozen_names).as_str(),
                        ))
                        .add(mutable_names, frozen_names, frames, last_time),
                },
            }
        }
    }

    fn render_with_buffer(&self, node: &mut FlameGraphNode) {
        let node = node.child(ArcStr::from(self.name));
        let count = self.time.to_duration().as_millis();
        if count > 0 {
            node.add(count as u64);
        }
        for x in self.children.values() {
            x.render_with_buffer(node);
        }
    }

    fn render(&self) -> FlameGraphData {
        let mut data = FlameGraphData::default();
        self.render_with_buffer(data.root());
        data
    }
}

impl<'v> TimeFlameProfile<'v> {
    pub(crate) fn new() -> Self {
        Self(None)
    }

    pub(crate) fn enable(&mut self) {
        self.0 = Some(Box::default());
    }

    pub(crate) fn record_call_enter(&mut self, function: Value<'v>) {
        if let Some(x) = &mut self.0 {
            let ind = x.index.index(function);
            x.frames.push((Frame::Push(ind), Instant::now()))
        }
    }

    pub(crate) fn record_call_exit(&mut self) {
        if let Some(x) = &mut self.0 {
            x.frames.push((Frame::Pop, Instant::now()))
        }
    }

    // We could expose profile on the Heap, but it's an implementation detail that it works here.
    pub(crate) fn gen(&self) -> anyhow::Result<ProfileData> {
        match &self.0 {
            None => Err(FlameProfileError::NotEnabled.into()),
            Some(x) => Ok(Self::gen_profile(x)),
        }
    }

    fn gen_profile(x: &FlameData) -> ProfileData {
        // Need to write out lines which look like:
        // root;calls1;calls2 1
        // All the numbers at the end must be whole numbers (we use milliseconds)
        let mutable_names = x.index.mutable_values.map(|x| x.to_repr());
        let frozen_names = x.index.frozen_values.map(|x| x.to_value().to_repr());
        ProfileData {
            profile_mode: ProfileMode::TimeFlame,
            profile: ProfileDataImpl::TimeFlameProfile(
                Stacks::new(&mutable_names, &frozen_names, &x.frames).render(),
            ),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::thread;
    use std::time::Duration;

    use anyhow::Context;
    use starlark_derive::starlark_module;

    use crate as starlark;
    use crate::assert::Assert;
    use crate::environment::Globals;
    use crate::environment::GlobalsBuilder;
    use crate::environment::Module;
    use crate::eval::runtime::file_loader::ReturnOwnedFileLoader;
    use crate::eval::Evaluator;
    use crate::eval::ProfileMode;
    use crate::syntax::AstModule;
    use crate::syntax::Dialect;
    use crate::values::none::NoneType;

    #[test]
    fn test_time_flame_works_inside_frozen_module() {
        #[starlark_module]
        fn register_sleep(globals: &mut GlobalsBuilder) {
            fn sleep() -> anyhow::Result<NoneType> {
                thread::sleep(Duration::from_millis(2));
                Ok(NoneType)
            }
        }

        let mut a = Assert::new();
        a.globals_add(register_sleep);
        let a_bzl = a.pass_module(
            r#"
def foo():
    for i in range(5):
        # Must sleep otherwise time flame will round the duration to zero and erase it.
        sleep()
    "#,
        );

        let modules = HashMap::from_iter([("a.bzl".to_owned(), a_bzl)]);
        let loader = ReturnOwnedFileLoader { modules };

        let module = Module::new();
        let mut eval = Evaluator::new(&module);
        eval.enable_profile(&ProfileMode::TimeFlame).unwrap();
        eval.set_loader(&loader);
        eval.eval_module(
            AstModule::parse(
                "x.star",
                r#"
load("a.bzl", "foo")

def bar():
    for i in range(10):
        foo()

bar()
"#
                .to_owned(),
                &Dialect::Standard,
            )
            .unwrap(),
            &Globals::standard(),
        )
        .unwrap();

        let profile = eval.gen_profile().unwrap().gen().unwrap();
        let the_line = profile
            .lines()
            .find(|l| l.contains("foo"))
            .with_context(|| {
                format!(
                    "There must be a line with `foo` in the profile: {:?}",
                    profile
                )
            })
            .unwrap();
        assert!(
            the_line.contains("bar"),
            "Profile must contain a line `bar.*foo`: {:?}",
            profile
        );
    }
}


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

//! Runtime typecheck profile.

use std::time::Duration;

use crate::collections::SmallMap;
use crate::eval::runtime::profile::csv::CsvWriter;
use crate::eval::runtime::profile::data::ProfileData;
use crate::eval::runtime::small_duration::SmallDuration;
use crate::eval::ProfileMode;
use crate::values::FrozenStringValue;

#[derive(Debug, thiserror::Error)]
enum TypecheckProfileError {
    #[error("Typecheck profile not enabled")]
    NotEnabled,
}

#[derive(Default, Debug)]
pub(crate) struct TypecheckProfile {
    pub(crate) enabled: bool,
    // TODO(nga): we don't need ordered map here.
    by_function: SmallMap<FrozenStringValue, SmallDuration>,
}

impl TypecheckProfile {
    pub(crate) fn add(&mut self, function: FrozenStringValue, time: Duration) {
        assert!(self.enabled);
        *self
            .by_function
            .entry_hashed(function.get_hashed())
            .or_default() += time;
    }

    fn gen_csv(&self) -> String {
        let total_time = self.by_function.values().sum::<SmallDuration>();

        let mut w = CsvWriter::new(["Function", "Time (s)"]);
        w.write_display("TOTAL");
        w.write_value(total_time);
        w.finish_row();

        let mut by_function = Vec::from_iter(&self.by_function);
        by_function.sort_by_key(|(name, t)| (u64::MAX - t.nanos, *name));

        for (name, t) in by_function {
            w.write_display(name.as_str());
            w.write_value(t);
            w.finish_row();
        }

        w.finish()
    }

    pub(crate) fn gen(&self) -> anyhow::Result<ProfileData> {
        if !self.enabled {
            return Err(TypecheckProfileError::NotEnabled.into());
        }
        Ok(ProfileData::new(ProfileMode::Typecheck, self.gen_csv()))
    }
}

#[cfg(test)]
mod tests {
    use crate::environment::Globals;
    use crate::environment::Module;
    use crate::eval::Evaluator;
    use crate::eval::ProfileMode;
    use crate::syntax::AstModule;
    use crate::syntax::Dialect;

    #[test]
    fn test_typecheck_profile() -> anyhow::Result<()> {
        let module = Module::new();
        let mut eval = Evaluator::new(&module);
        let program = r#"
def f(s: str.type):
    return int(s)

def g():
    for i in range(0, 1000):
        f(str(i))

g()
"#;
        let program = AstModule::parse("test.star", program.to_owned(), &Dialect::Extended)?;
        eval.enable_profile(&ProfileMode::Typecheck)?;
        eval.eval_module(program, &Globals::extended())?;

        let csv = eval.typecheck_profile.gen_csv();
        let lines: Vec<&str> = csv.lines().collect();
        assert_eq!("Function,Time (s)", lines[0]);
        assert!(lines[1].starts_with("\"TOTAL\","), "{:?}", lines[1]);
        assert!(lines[2].starts_with("\"f\","), "{:?}", lines[2]);
        assert_eq!(3, lines.len());

        Ok(())
    }
}


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

//! The server that allows IDEs to evaluate and interpret starlark code according
//! to the [Language Server Protocol](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/).

pub mod server;
mod symbols;
#[cfg(all(test, not(windows)))]
mod test;