sightline-ty-python-core 0.1.0

use bitflags::bitflags;
use hashbrown::hash_table::Entry;
use ruff_index::{IndexVec, newtype_index};
use ruff_python_ast::name::Name;
use rustc_hash::FxHasher;
use std::hash::{Hash as _, Hasher as _};
use std::ops::{Deref, DerefMut};

// Selected using performance and memory profiling across the 162-project ecosystem corpus.
// Symbol-name equality is cheap enough that raising the cutoff from 8 to 16 reduced retained
// memory without a measurable performance regression.
const LINEAR_SEARCH_THRESHOLD: usize = 16;

/// Uniquely identifies a symbol in a given scope.
#[newtype_index]
#[derive(Ord, PartialOrd, get_size2::GetSize)]
pub struct ScopedSymbolId;

/// A symbol in a given scope.
#[derive(Debug, Clone, PartialEq, Eq, get_size2::GetSize)]
pub struct Symbol {
    name: Name,
    flags: SymbolFlags,
}

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

bitflags! {
    /// Flags that can be queried to obtain information about a symbol in a given scope.
    ///
    /// See the doc-comment at the top of [`super::use_def`] for explanations of what it
    /// means for a symbol to be *bound* as opposed to *declared*.
    #[derive(Copy, Clone, Debug, Eq, PartialEq)]
    struct SymbolFlags: u8 {
        const IS_USED               = 1 << 0;
        const IS_BOUND              = 1 << 1;
        const IS_DECLARED           = 1 << 2;
        const MARKED_GLOBAL         = 1 << 3;
        const MARKED_NONLOCAL       = 1 << 4;
        /// true if the symbol is assigned more than once, or if it is assigned even though it is already in use
        const IS_REASSIGNED         = 1 << 5;
        const IS_PARAMETER          = 1 << 6;
    }
}

impl get_size2::GetSize for SymbolFlags {}

impl Symbol {
    pub(crate) const fn new(name: Name) -> Self {
        Self {
            name,
            flags: SymbolFlags::empty(),
        }
    }

    pub fn name(&self) -> &Name {
        &self.name
    }

    /// Is the symbol used in its containing scope?
    pub fn is_used(&self) -> bool {
        self.flags.contains(SymbolFlags::IS_USED)
    }

    /// Is the symbol given a value in its containing scope?
    pub const fn is_bound(&self) -> bool {
        self.flags.contains(SymbolFlags::IS_BOUND)
    }

    /// Is the symbol declared in its containing scope?
    pub fn is_declared(&self) -> bool {
        self.flags.contains(SymbolFlags::IS_DECLARED)
    }

    /// Is the symbol `global` its containing scope?
    pub fn is_global(&self) -> bool {
        self.flags.contains(SymbolFlags::MARKED_GLOBAL)
    }

    /// Is the symbol `nonlocal` its containing scope?
    pub fn is_nonlocal(&self) -> bool {
        self.flags.contains(SymbolFlags::MARKED_NONLOCAL)
    }

    /// Is the symbol defined in this scope, vs referring to some enclosing scope?
    ///
    /// There are three common cases where a name refers to an enclosing scope:
    ///
    /// 1. explicit `global` variables
    /// 2. explicit `nonlocal` variables
    /// 3. "free" variables, which are used in a scope where they're neither bound nor declared
    ///
    /// Note that even if `is_local` is false, that doesn't necessarily mean there's an enclosing
    /// scope that resolves the reference. The symbol could be a built-in like `print`, or a name
    /// error at runtime, or a global variable added dynamically with e.g. `globals()`.
    ///
    /// XXX: There's a fourth case that we don't (can't) handle here. A variable that's bound or
    /// declared (anywhere) in a class body, but used before it's bound (at runtime), resolves
    /// (unbelievably) to the global scope. For example:
    /// ```py
    /// x = 42
    /// def f():
    ///     x = 43
    ///     class Foo:
    ///         print(x)  # 42 (never 43)
    ///         if secrets.randbelow(2):
    ///             x = 44
    ///         print(x)  # 42 or 44
    /// ```
    /// In cases like this, the resolution isn't known until runtime, and in fact it varies from
    /// one use to the next. The semantic index alone can't resolve this, and instead it's a
    /// special case in type inference (see `infer_place_load`).
    pub fn is_local(&self) -> bool {
        !self.is_global() && !self.is_nonlocal() && (self.is_bound() || self.is_declared())
    }

    pub const fn is_reassigned(&self) -> bool {
        self.flags.contains(SymbolFlags::IS_REASSIGNED)
    }

    pub(crate) fn is_parameter(&self) -> bool {
        self.flags.contains(SymbolFlags::IS_PARAMETER)
    }

    pub(super) fn mark_global(&mut self) {
        self.insert_flags(SymbolFlags::MARKED_GLOBAL);
    }

    pub(super) fn mark_nonlocal(&mut self) {
        self.insert_flags(SymbolFlags::MARKED_NONLOCAL);
    }

    pub(super) fn mark_bound(&mut self) {
        if self.is_bound() || self.is_used() {
            self.insert_flags(SymbolFlags::IS_REASSIGNED);
        }

        self.insert_flags(SymbolFlags::IS_BOUND);
    }

    pub(super) fn mark_used(&mut self) {
        self.insert_flags(SymbolFlags::IS_USED);
    }

    pub(super) fn mark_declared(&mut self) {
        self.insert_flags(SymbolFlags::IS_DECLARED);
    }

    pub(super) fn mark_parameter(&mut self) {
        self.insert_flags(SymbolFlags::IS_PARAMETER);
    }

    fn insert_flags(&mut self, flags: SymbolFlags) {
        self.flags.insert(flags);
    }
}

/// Map from symbol name to its ID.
///
/// Uses a hash table to avoid storing the name twice.
#[derive(Debug, Default, get_size2::GetSize)]
struct SymbolReverseTable(hashbrown::HashTable<ScopedSymbolId>);

impl SymbolReverseTable {
    fn symbol_id(
        &self,
        symbols: &IndexVec<ScopedSymbolId, Symbol>,
        name: &str,
    ) -> Option<ScopedSymbolId> {
        self.0
            .find(Self::hash_name(name), |id| symbols[*id].name == name)
            .copied()
    }

    fn entry<'a>(
        &'a mut self,
        symbols: &IndexVec<ScopedSymbolId, Symbol>,
        symbol: &Symbol,
    ) -> Entry<'a, ScopedSymbolId> {
        self.0.entry(
            Self::hash_name(symbol.name()),
            |id| &symbols[*id].name == symbol.name(),
            |id| Self::hash_name(&symbols[*id].name),
        )
    }

    fn shrink_to_fit(&mut self, symbols: &IndexVec<ScopedSymbolId, Symbol>) {
        self.0
            .shrink_to_fit(|id| Self::hash_name(&symbols[*id].name));
    }

    fn hash_name(name: &str) -> u64 {
        let mut h = FxHasher::default();
        name.hash(&mut h);
        h.finish()
    }
}

/// The symbols of a given scope.
///
/// Allows lookup by name and a symbol's ID.
#[derive(Default, get_size2::GetSize)]
pub(super) struct SymbolTable {
    symbols: IndexVec<ScopedSymbolId, Symbol>,
    /// Reverse lookup retained only when linear search would be expensive.
    reverse: Option<Box<SymbolReverseTable>>,
}

impl SymbolTable {
    /// Look up a symbol by its ID.
    ///
    /// ## Panics
    /// If the ID is not valid for this symbol table.
    #[track_caller]
    pub(crate) fn symbol(&self, id: ScopedSymbolId) -> &Symbol {
        &self.symbols[id]
    }

    /// Look up a symbol by its ID, mutably.
    ///
    /// ## Panics
    /// If the ID is not valid for this symbol table.
    #[track_caller]
    pub(crate) fn symbol_mut(&mut self, id: ScopedSymbolId) -> &mut Symbol {
        &mut self.symbols[id]
    }

    /// Look up the ID of a symbol by its name.
    pub(crate) fn symbol_id(&self, name: &str) -> Option<ScopedSymbolId> {
        if let Some(reverse) = self.reverse.as_deref() {
            return reverse.symbol_id(&self.symbols, name);
        }

        self.symbols
            .iter_enumerated()
            .find_map(|(id, symbol)| (symbol.name == name).then_some(id))
    }

    /// Iterate over the symbols in this symbol table.
    pub(crate) fn iter(&self) -> std::slice::Iter<'_, Symbol> {
        self.symbols.iter()
    }
}

impl PartialEq for SymbolTable {
    fn eq(&self, other: &Self) -> bool {
        // It's sufficient to compare the symbols as the map is only a reverse lookup.
        self.symbols == other.symbols
    }
}

impl Eq for SymbolTable {}

impl std::fmt::Debug for SymbolTable {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_tuple("SymbolTable").field(&self.symbols).finish()
    }
}

#[derive(Debug, Default)]
pub(super) struct SymbolTableBuilder {
    table: SymbolTable,
    reverse: SymbolReverseTable,
}

impl SymbolTableBuilder {
    pub(super) fn symbol_id(&self, name: &str) -> Option<ScopedSymbolId> {
        self.reverse.symbol_id(&self.table.symbols, name)
    }

    /// Add a new symbol to this scope or update the flags if a symbol with the same name already exists.
    pub(super) fn add(&mut self, symbol: Symbol) -> (ScopedSymbolId, bool) {
        let entry = self.reverse.entry(&self.table.symbols, &symbol);

        match entry {
            Entry::Occupied(entry) => {
                let id = *entry.get();

                if !symbol.flags.is_empty() {
                    self.symbols[id].flags.insert(symbol.flags);
                }

                (id, false)
            }
            Entry::Vacant(entry) => {
                let id = self.table.symbols.push(symbol);
                entry.insert(id);
                (id, true)
            }
        }
    }

    pub(super) fn build(self) -> SymbolTable {
        let Self {
            mut table,
            mut reverse,
        } = self;
        table.symbols.shrink_to_fit();

        if table.symbols.len() > LINEAR_SEARCH_THRESHOLD {
            reverse.shrink_to_fit(&table.symbols);
            table.reverse = Some(Box::new(reverse));
        }

        table
    }
}

impl Deref for SymbolTableBuilder {
    type Target = SymbolTable;

    fn deref(&self) -> &Self::Target {
        &self.table
    }
}

impl DerefMut for SymbolTableBuilder {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.table
    }
}