sway-types 0.71.0

use crate::{span, LineCol, ProgramId, SourceId, Span};
use parking_lot::RwLock;
use std::{
    collections::{BTreeSet, HashMap},
    path::{Path, PathBuf},
    str::FromStr,
};
use toml::Table;

/// The Source Engine manages a relationship between file paths and their corresponding
/// integer-based source IDs. Additionally, it maintains the reverse - a map that traces
/// back from a source ID to its original file path. The primary objective of this
/// system is to enable clients that need to reference a file path to do so using an
/// integer-based ID. This numeric representation can be stored more efficiently as
/// a key in a hashmap.
/// The Source Engine is designed to be thread-safe. Its internal structures are
/// secured by the RwLock mechanism. This allows its functions to be invoked using
/// a straightforward non-mutable reference, ensuring safe concurrent access.
#[derive(Debug, Default)]
pub struct SourceEngine {
    next_source_id: RwLock<u32>,
    path_to_source_map: RwLock<HashMap<PathBuf, SourceId>>,
    source_to_path_map: RwLock<HashMap<SourceId, PathBuf>>,
    source_to_buffer_map: RwLock<HashMap<SourceId, span::Source>>,
    next_program_id: RwLock<u16>,
    manifest_path_to_program_map: RwLock<HashMap<PathBuf, ProgramId>>,
    /// Stores the package name and version for manifest path,
    /// if available in the Forc.toml file, or the fallback package name
    /// coming from the last part of the manifest path, and an empty version string.
    manifest_path_to_package_info: RwLock<HashMap<PathBuf, (String, String)>>,
    module_to_sources_map: RwLock<HashMap<ProgramId, BTreeSet<SourceId>>>,
}

impl Clone for SourceEngine {
    fn clone(&self) -> Self {
        SourceEngine {
            next_source_id: RwLock::new(*self.next_source_id.read()),
            path_to_source_map: RwLock::new(self.path_to_source_map.read().clone()),
            source_to_path_map: RwLock::new(self.source_to_path_map.read().clone()),
            source_to_buffer_map: RwLock::new(self.source_to_buffer_map.read().clone()),
            next_program_id: RwLock::new(*self.next_program_id.read()),
            manifest_path_to_program_map: RwLock::new(
                self.manifest_path_to_program_map.read().clone(),
            ),
            manifest_path_to_package_info: RwLock::new(
                self.manifest_path_to_package_info.read().clone(),
            ),
            module_to_sources_map: RwLock::new(self.module_to_sources_map.read().clone()),
        }
    }
}

impl SourceEngine {
    const AUTOGENERATED_PATH: &'static str = "<autogenerated>";

    pub fn is_span_in_autogenerated(&self, span: &crate::Span) -> Option<bool> {
        span.source_id().map(|s| self.is_source_id_autogenerated(s))
    }

    pub fn is_source_id_autogenerated(&self, source_id: &SourceId) -> bool {
        self.get_path(source_id)
            .display()
            .to_string()
            .contains("<autogenerated>")
    }

    /// Fetch the cached source buffer for `source_id`, replacing it if the caller
    /// provides newer text so future span lookups observe the latest file contents.
    pub fn get_or_create_source_buffer(
        &self,
        source_id: &SourceId,
        source: span::Source,
    ) -> span::Source {
        let mut map = self.source_to_buffer_map.write();

        if let Some(existing) = map.get_mut(source_id) {
            // Replace the cached source if the contents have changed so that
            // subsequent spans reflect the latest file text.
            if existing.text != source.text {
                *existing = source;
            }
            existing.clone()
        } else {
            map.insert(*source_id, source.clone());
            source
        }
    }

    /// This function retrieves an integer-based source ID for a provided path buffer.
    /// If an ID already exists for the given path, the function will return that
    /// existing ID. If not, a new ID will be created.
    pub fn get_source_id(&self, path: &PathBuf) -> SourceId {
        {
            let source_map = self.path_to_source_map.read();
            if source_map.contains_key(path) {
                return source_map.get(path).copied().unwrap();
            }
        }

        let program_id = self.get_or_create_program_id_from_manifest_path(path);
        self.get_source_id_with_program_id(path, program_id)
    }

    pub fn get_source_id_with_program_id(&self, path: &PathBuf, program_id: ProgramId) -> SourceId {
        {
            let source_map = self.path_to_source_map.read();
            if source_map.contains_key(path) {
                return source_map.get(path).copied().unwrap();
            }
        }

        let source_id = SourceId::new(program_id.0, *self.next_source_id.read());
        {
            let mut next_id = self.next_source_id.write();
            *next_id += 1;

            let mut source_map = self.path_to_source_map.write();
            source_map.insert(path.clone(), source_id);

            let mut path_map = self.source_to_path_map.write();
            path_map.insert(source_id, path.clone());
        }

        let mut module_map = self.module_to_sources_map.write();
        module_map.entry(program_id).or_default().insert(source_id);

        source_id
    }

    /// Return the associated autogenerated pseudo file for the passed `source_id`.
    /// Example: main.autogenerated.sw for main.sw
    ///
    /// Returns `None`, if `source_id` does not have a valid path.
    pub fn get_associated_autogenerated_source_id(&self, source_id: &SourceId) -> Option<SourceId> {
        let path = self.get_path(source_id);
        let file_name = PathBuf::from_str(path.file_name()?.to_str()?).ok()?;
        let path = path.with_file_name(format!(
            "{}.{}.{}",
            file_name.file_stem()?.to_str()?,
            Self::AUTOGENERATED_PATH,
            file_name.extension()?.to_str()?
        ));
        Some(self.get_source_id_with_program_id(&path, source_id.program_id()))
    }

    /// This function provides the file path corresponding to a specified source ID.
    pub fn get_path(&self, source_id: &SourceId) -> PathBuf {
        self.source_to_path_map
            .read()
            .get(source_id)
            .unwrap()
            .clone()
    }

    /// This function provides the [ProgramId] corresponding to a specified manifest file path.
    pub fn get_program_id_from_manifest_path(&self, path: impl AsRef<Path>) -> Option<ProgramId> {
        let manifest_path = sway_utils::find_parent_manifest_dir(&path)
            .unwrap_or_else(|| path.as_ref().to_path_buf());
        self.manifest_path_to_program_map
            .read()
            .get(&manifest_path)
            .copied()
    }

    pub fn get_or_create_program_id_from_manifest_path(&self, path: &PathBuf) -> ProgramId {
        let manifest_path = sway_utils::find_parent_manifest_dir(path).unwrap_or(path.clone());
        let mut module_map = self.manifest_path_to_program_map.write();
        *module_map.entry(manifest_path.clone()).or_insert_with(|| {
            let mut next_id = self.next_program_id.write();
            *next_id += 1;
            ProgramId::new(*next_id)
        })
    }

    /// Returns the [PathBuf] associated with the provided [ProgramId], if it exists in the manifest_path_to_program_map.
    pub fn get_manifest_path_from_program_id(&self, program_id: &ProgramId) -> Option<PathBuf> {
        let path_to_module_map = self.manifest_path_to_program_map.read();
        path_to_module_map
            .iter()
            .find(|(_, &id)| id == *program_id)
            .map(|(path, _)| path.clone())
    }

    /// This function provides the file name (with extension) corresponding to a specified source ID.
    pub fn get_file_name(&self, source_id: &SourceId) -> Option<String> {
        self.get_path(source_id)
            .as_path()
            .file_name()
            .map(|file_name| file_name.to_string_lossy())
            .map(|file_name| file_name.to_string())
    }

    pub fn all_files(&self) -> Vec<PathBuf> {
        let s = self.source_to_path_map.read();
        let mut v = s.values().cloned().collect::<Vec<_>>();
        v.sort();
        v
    }

    pub fn get_source_ids_from_program_id(
        &self,
        program_id: ProgramId,
    ) -> Option<BTreeSet<SourceId>> {
        let s = self.module_to_sources_map.read();
        s.get(&program_id).cloned()
    }

    // TODO: Do we want to parse Forc.toml files here, within the `SourceEngine`?
    //       Currently, we don't have any other option to obtain the original package name and version.
    //       But ideally, this and potentially other information should be passed to the compiler
    //       from the tooling layer.
    fn get_package_name_and_version(&self, manifest_path: &Path) -> (String, String) {
        fn get_fallback_package_name_and_version(manifest_path: &Path) -> (String, String) {
            // As a fallback, use the last part of the manifest path as the package name.
            // This should actually never happen, because we should always have a valid
            // Forc.toml file at the manifest path.
            let package_dir_name = manifest_path
                .iter()
                .next_back()
                .map(|p| p.to_string_lossy().to_string())
                .unwrap_or_else(|| "<unknown>".to_string());
            (package_dir_name, String::new())
        }

        fn get_project_field(toml: &Table, field_name: &str) -> Option<String> {
            toml.get("project")
                .and_then(|v| v.get(field_name))
                .and_then(|field| field.as_str())
                .map(|value| value.to_string())
        }

        let forc_toml_path = manifest_path.join("Forc.toml");
        if !forc_toml_path.exists() {
            return get_fallback_package_name_and_version(manifest_path);
        }

        let content = match std::fs::read_to_string(&forc_toml_path) {
            Ok(content) => content,
            Err(_) => return get_fallback_package_name_and_version(manifest_path),
        };

        let toml = match content.parse::<Table>() {
            Ok(toml) => toml,
            Err(_) => return get_fallback_package_name_and_version(manifest_path),
        };

        let package_name = get_project_field(&toml, "name").unwrap_or("<unknown>".to_string());
        let package_version = get_project_field(&toml, "version").unwrap_or_default();

        (package_name, package_version)
    }

    pub fn get_source_location(&self, span: &Span) -> SourceLocation {
        let Some(source_id) = span.source_id() else {
            return SourceLocation::unknown();
        };

        let source_file = self.get_path(source_id);

        // Find the manifest path from the source file path.
        let program_id = self
            .get_program_id_from_manifest_path(&source_file)
            .expect("the `source_file` is retrieved from the `SourceEngine::get_path` function so the manifest path and program id must exist");
        let manifest_path = self.get_manifest_path_from_program_id(&program_id).expect(
            "the `program_id` is retrieved from the `SourceEngine` so the manifest path must exist",
        );

        // TODO: Use HashSet::get_or_insert_with once it gets available (currently experimental).
        //       See: https://doc.rust-lang.org/std/collections/struct.HashSet.html#method.get_or_insert_with
        let mut package_infos = self.manifest_path_to_package_info.write();
        let (package_name, package_version) = &package_infos
            .entry(manifest_path.clone())
            .or_insert_with(|| self.get_package_name_and_version(&manifest_path));

        let pkg = if package_version.is_empty() {
            package_name.clone()
        } else {
            format!("{package_name}@{package_version}")
        };

        // Get the relative path of the source file with respect to the manifest path.
        let source_file = source_file
            .strip_prefix(&manifest_path)
            .expect("the `manifest_path` is a parent of the `source_file`")
            .to_string_lossy();
        let source_file = if let Some(source_file_no_leading_slash) = source_file.strip_prefix("/")
        {
            source_file_no_leading_slash.to_string()
        } else {
            source_file.to_string()
        };

        SourceLocation {
            pkg,
            file: source_file,
            loc: span.start_line_col_one_index(),
        }
    }
}

/// A location in a source code file, represented by a package name, relative file path,
/// and line/column within the file.
#[derive(Default, Debug, Clone, PartialEq, Eq, Hash)]
pub struct SourceLocation {
    /// The name and the version of the package that contains the source file,
    /// in the format "name@version".
    ///
    /// E.g., "my_lib@0.1.0".
    ///
    /// The version is optional and may be omitted. E.g., "my_lib".
    pub pkg: String,
    /// The path to the source file, relative to the package root.
    ///
    /// E.g., "src/lib.rs".
    pub file: String,
    pub loc: LineCol,
}

impl SourceLocation {
    /// Creates a new unknown `SourceLocation` instance, which has
    /// package and file set to "<unknown>".
    pub fn unknown() -> Self {
        Self {
            pkg: "<unknown>".to_string(),
            file: "<unknown>".to_string(),
            loc: LineCol::default(),
        }
    }
}