compile.rs

//!
//! Take an AST and transform it into bytecode
//!
//! Inspirational code:
//!   <https://github.com/python/cpython/blob/main/Python/compile.c>
//!   <https://github.com/micropython/micropython/blob/master/py/compile.c>

// spell-checker:ignore starunpack subscripter

#![deny(clippy::cast_possible_truncation)]

use crate::{
    IndexMap, IndexSet, ToPythonName,
    error::{CodegenError, CodegenErrorType, InternalError, PatternUnreachableReason},
    ir::{self, BlockIdx},
    symboltable::{self, CompilerScope, SymbolFlags, SymbolScope, SymbolTable},
    unparse::UnparseExpr,
};
use itertools::Itertools;
use malachite_bigint::BigInt;
use num_complex::Complex;
use num_traits::{Num, ToPrimitive};
use ruff_python_ast::{
    Alias, Arguments, BoolOp, CmpOp, Comprehension, ConversionFlag, DebugText, Decorator, DictItem,
    ExceptHandler, ExceptHandlerExceptHandler, Expr, ExprAttribute, ExprBoolOp, ExprContext,
    ExprFString, ExprList, ExprName, ExprSlice, ExprStarred, ExprSubscript, ExprTuple, ExprUnaryOp,
    FString, FStringFlags, FStringPart, Identifier, Int, InterpolatedElement,
    InterpolatedStringElement, InterpolatedStringElements, Keyword, MatchCase, ModExpression,
    ModModule, Operator, Parameters, Pattern, PatternMatchAs, PatternMatchClass,
    PatternMatchMapping, PatternMatchOr, PatternMatchSequence, PatternMatchSingleton,
    PatternMatchStar, PatternMatchValue, Singleton, Stmt, StmtExpr, TypeParam, TypeParamParamSpec,
    TypeParamTypeVar, TypeParamTypeVarTuple, TypeParams, UnaryOp, WithItem,
};
use ruff_text_size::{Ranged, TextRange};
use rustpython_compiler_core::{
    Mode, OneIndexed, PositionEncoding, SourceFile, SourceLocation,
    bytecode::{
        self, Arg as OpArgMarker, BinaryOperator, CodeObject, ComparisonOperator, ConstantData,
        Instruction, OpArg, OpArgType, UnpackExArgs,
    },
};
use rustpython_wtf8::Wtf8Buf;
use std::{borrow::Cow, collections::HashSet};

const MAXBLOCKS: usize = 20;

#[derive(Debug, Clone, Copy)]
pub enum FBlockType {
    WhileLoop,
    ForLoop,
    TryExcept,
    FinallyTry,
    FinallyEnd,
    With,
    AsyncWith,
    HandlerCleanup,
    PopValue,
    ExceptionHandler,
    ExceptionGroupHandler,
    AsyncComprehensionGenerator,
    StopIteration,
}

#[derive(Debug, Clone)]
pub struct FBlockInfo {
    pub fb_type: FBlockType,
    pub fb_block: BlockIdx,
    pub fb_exit: BlockIdx,
    // fb_datum is not needed in RustPython
}

pub(crate) type InternalResult<T> = Result<T, InternalError>;
type CompileResult<T> = Result<T, CodegenError>;

#[derive(PartialEq, Eq, Clone, Copy)]
enum NameUsage {
    Load,
    Store,
    Delete,
}

enum CallType {
    Positional { nargs: u32 },
    Keyword { nargs: u32 },
    Ex { has_kwargs: bool },
}

fn is_forbidden_name(name: &str) -> bool {
    // See https://docs.python.org/3/library/constants.html#built-in-constants
    const BUILTIN_CONSTANTS: &[&str] = &["__debug__"];

    BUILTIN_CONSTANTS.contains(&name)
}

/// Main structure holding the state of compilation.
struct Compiler {
    code_stack: Vec<ir::CodeInfo>,
    symbol_table_stack: Vec<SymbolTable>,
    source_file: SourceFile,
    // current_source_location: SourceLocation,
    current_source_range: TextRange,
    done_with_future_stmts: DoneWithFuture,
    future_annotations: bool,
    ctx: CompileContext,
    opts: CompileOpts,
    in_annotation: bool,
}

enum DoneWithFuture {
    No,
    DoneWithDoc,
    Yes,
}

#[derive(Debug, Clone, Default)]
pub struct CompileOpts {
    /// How optimized the bytecode output should be; any optimize > 0 does
    /// not emit assert statements
    pub optimize: u8,
}

#[derive(Debug, Clone, Copy)]
struct CompileContext {
    loop_data: Option<(ir::BlockIdx, ir::BlockIdx)>,
    in_class: bool,
    func: FunctionContext,
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum FunctionContext {
    NoFunction,
    Function,
    AsyncFunction,
}

impl CompileContext {
    fn in_func(self) -> bool {
        self.func != FunctionContext::NoFunction
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum ComprehensionType {
    Generator,
    List,
    Set,
    Dict,
}

fn validate_duplicate_params(params: &Parameters) -> Result<(), CodegenErrorType> {
    let mut seen_params = HashSet::new();
    for param in params {
        let param_name = param.name().as_str();
        if !seen_params.insert(param_name) {
            return Err(CodegenErrorType::SyntaxError(format!(
                r#"Duplicate parameter "{param_name}""#
            )));
        }
    }

    Ok(())
}

/// Compile an Mod produced from ruff parser
pub fn compile_top(
    ast: ruff_python_ast::Mod,
    source_file: SourceFile,
    mode: Mode,
    opts: CompileOpts,
) -> CompileResult<CodeObject> {
    match ast {
        ruff_python_ast::Mod::Module(module) => match mode {
            Mode::Exec | Mode::Eval => compile_program(&module, source_file, opts),
            Mode::Single => compile_program_single(&module, source_file, opts),
            Mode::BlockExpr => compile_block_expression(&module, source_file, opts),
        },
        ruff_python_ast::Mod::Expression(expr) => compile_expression(&expr, source_file, opts),
    }
}

/// Compile a standard Python program to bytecode
pub fn compile_program(
    ast: &ModModule,
    source_file: SourceFile,
    opts: CompileOpts,
) -> CompileResult<CodeObject> {
    let symbol_table = SymbolTable::scan_program(ast, source_file.clone())
        .map_err(|e| e.into_codegen_error(source_file.name().to_owned()))?;
    let mut compiler = Compiler::new(opts, source_file, "<module>".to_owned());
    compiler.compile_program(ast, symbol_table)?;
    let code = compiler.exit_scope();
    trace!("Compilation completed: {code:?}");
    Ok(code)
}

/// Compile a Python program to bytecode for the context of a REPL
pub fn compile_program_single(
    ast: &ModModule,
    source_file: SourceFile,
    opts: CompileOpts,
) -> CompileResult<CodeObject> {
    let symbol_table = SymbolTable::scan_program(ast, source_file.clone())
        .map_err(|e| e.into_codegen_error(source_file.name().to_owned()))?;
    let mut compiler = Compiler::new(opts, source_file, "<module>".to_owned());
    compiler.compile_program_single(&ast.body, symbol_table)?;
    let code = compiler.exit_scope();
    trace!("Compilation completed: {code:?}");
    Ok(code)
}

pub fn compile_block_expression(
    ast: &ModModule,
    source_file: SourceFile,
    opts: CompileOpts,
) -> CompileResult<CodeObject> {
    let symbol_table = SymbolTable::scan_program(ast, source_file.clone())
        .map_err(|e| e.into_codegen_error(source_file.name().to_owned()))?;
    let mut compiler = Compiler::new(opts, source_file, "<module>".to_owned());
    compiler.compile_block_expr(&ast.body, symbol_table)?;
    let code = compiler.exit_scope();
    trace!("Compilation completed: {code:?}");
    Ok(code)
}

pub fn compile_expression(
    ast: &ModExpression,
    source_file: SourceFile,
    opts: CompileOpts,
) -> CompileResult<CodeObject> {
    let symbol_table = SymbolTable::scan_expr(ast, source_file.clone())
        .map_err(|e| e.into_codegen_error(source_file.name().to_owned()))?;
    let mut compiler = Compiler::new(opts, source_file, "<module>".to_owned());
    compiler.compile_eval(ast, symbol_table)?;
    let code = compiler.exit_scope();
    Ok(code)
}

macro_rules! emit {
    ($c:expr, Instruction::$op:ident { $arg:ident$(,)? }$(,)?) => {
        $c.emit_arg($arg, |x| Instruction::$op { $arg: x })
    };
    ($c:expr, Instruction::$op:ident { $arg:ident : $arg_val:expr $(,)? }$(,)?) => {
        $c.emit_arg($arg_val, |x| Instruction::$op { $arg: x })
    };
    ($c:expr, Instruction::$op:ident( $arg_val:expr $(,)? )$(,)?) => {
        $c.emit_arg($arg_val, Instruction::$op)
    };
    ($c:expr, Instruction::$op:ident$(,)?) => {
        $c.emit_no_arg(Instruction::$op)
    };
}

fn eprint_location(zelf: &Compiler) {
    let start = zelf
        .source_file
        .to_source_code()
        .source_location(zelf.current_source_range.start(), PositionEncoding::Utf8);
    let end = zelf
        .source_file
        .to_source_code()
        .source_location(zelf.current_source_range.end(), PositionEncoding::Utf8);
    eprintln!(
        "LOCATION: {} from {}:{} to {}:{}",
        zelf.source_file.name(),
        start.line,
        start.character_offset,
        end.line,
        end.character_offset
    );
}

/// Better traceback for internal error
#[track_caller]
fn unwrap_internal<T>(zelf: &Compiler, r: InternalResult<T>) -> T {
    if let Err(ref r_err) = r {
        eprintln!("=== CODEGEN PANIC INFO ===");
        eprintln!("This IS an internal error: {r_err}");
        eprint_location(zelf);
        eprintln!("=== END PANIC INFO ===");
    }
    r.unwrap()
}

fn compiler_unwrap_option<T>(zelf: &Compiler, o: Option<T>) -> T {
    if o.is_none() {
        eprintln!("=== CODEGEN PANIC INFO ===");
        eprintln!("This IS an internal error, an option was unwrapped during codegen");
        eprint_location(zelf);
        eprintln!("=== END PANIC INFO ===");
    }
    o.unwrap()
}

// fn compiler_result_unwrap<T, E: std::fmt::Debug>(zelf: &Compiler, result: Result<T, E>) -> T {
//     if result.is_err() {
//         eprintln!("=== CODEGEN PANIC INFO ===");
//         eprintln!("This IS an internal error, an result was unwrapped during codegen");
//         eprint_location(zelf);
//         eprintln!("=== END PANIC INFO ===");
//     }
//     result.unwrap()
// }

/// The pattern context holds information about captured names and jump targets.
#[derive(Clone)]
pub struct PatternContext {
    /// A list of names captured by the pattern.
    pub stores: Vec<String>,
    /// If false, then any name captures against our subject will raise.
    pub allow_irrefutable: bool,
    /// A list of jump target labels used on pattern failure.
    pub fail_pop: Vec<BlockIdx>,
    /// The number of items on top of the stack that should remain.
    pub on_top: usize,
}

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

impl PatternContext {
    pub const fn new() -> Self {
        Self {
            stores: Vec::new(),
            allow_irrefutable: false,
            fail_pop: Vec::new(),
            on_top: 0,
        }
    }

    pub fn fail_pop_size(&self) -> usize {
        self.fail_pop.len()
    }
}

enum JumpOp {
    Jump,
    PopJumpIfFalse,
}

/// Type of collection to build in starunpack_helper
#[derive(Debug, Clone, Copy, PartialEq)]
enum CollectionType {
    Tuple,
    List,
    Set,
}

impl Compiler {
    fn new(opts: CompileOpts, source_file: SourceFile, code_name: String) -> Self {
        let module_code = ir::CodeInfo {
            flags: bytecode::CodeFlags::NEW_LOCALS,
            source_path: source_file.name().to_owned(),
            private: None,
            blocks: vec![ir::Block::default()],
            current_block: ir::BlockIdx(0),
            metadata: ir::CodeUnitMetadata {
                name: code_name.clone(),
                qualname: Some(code_name),
                consts: IndexSet::default(),
                names: IndexSet::default(),
                varnames: IndexSet::default(),
                cellvars: IndexSet::default(),
                freevars: IndexSet::default(),
                fast_hidden: IndexMap::default(),
                argcount: 0,
                posonlyargcount: 0,
                kwonlyargcount: 0,
                firstlineno: OneIndexed::MIN,
            },
            static_attributes: None,
            in_inlined_comp: false,
            fblock: Vec::with_capacity(MAXBLOCKS),
            symbol_table_index: 0, // Module is always the first symbol table
        };
        Self {
            code_stack: vec![module_code],
            symbol_table_stack: Vec::new(),
            source_file,
            // current_source_location: SourceLocation::default(),
            current_source_range: TextRange::default(),
            done_with_future_stmts: DoneWithFuture::No,
            future_annotations: false,
            ctx: CompileContext {
                loop_data: None,
                in_class: false,
                func: FunctionContext::NoFunction,
            },
            opts,
            in_annotation: false,
        }
    }

    /// Check if the slice is a two-element slice (no step)
    // = is_two_element_slice
    const fn is_two_element_slice(slice: &Expr) -> bool {
        matches!(slice, Expr::Slice(s) if s.step.is_none())
    }

    /// Compile a slice expression
    // = compiler_slice
    fn compile_slice(&mut self, s: &ExprSlice) -> CompileResult<u32> {
        // Compile lower
        if let Some(lower) = &s.lower {
            self.compile_expression(lower)?;
        } else {
            self.emit_load_const(ConstantData::None);
        }

        // Compile upper
        if let Some(upper) = &s.upper {
            self.compile_expression(upper)?;
        } else {
            self.emit_load_const(ConstantData::None);
        }

        // Compile step if present
        if let Some(step) = &s.step {
            self.compile_expression(step)?;
            Ok(3) // Three values on stack
        } else {
            Ok(2) // Two values on stack
        }
    }

    /// Compile a subscript expression
    // = compiler_subscript
    fn compile_subscript(
        &mut self,
        value: &Expr,
        slice: &Expr,
        ctx: ExprContext,
    ) -> CompileResult<()> {
        // 1. Check subscripter and index for Load context
        // 2. VISIT value
        // 3. Handle two-element slice specially
        // 4. Otherwise VISIT slice and emit appropriate instruction

        // For Load context, CPython does some checks (we skip for now)
        // if ctx == ExprContext::Load {
        //     check_subscripter(value);
        //     check_index(value, slice);
        // }

        // VISIT(c, expr, e->v.Subscript.value)
        self.compile_expression(value)?;

        // Handle two-element slice (for Load/Store, not Del)
        if Self::is_two_element_slice(slice) && !matches!(ctx, ExprContext::Del) {
            let n = match slice {
                Expr::Slice(s) => self.compile_slice(s)?,
                _ => unreachable!("is_two_element_slice should only return true for Expr::Slice"),
            };
            match ctx {
                ExprContext::Load => {
                    // CPython uses BINARY_SLICE
                    emit!(self, Instruction::BuildSlice { step: n == 3 });
                    emit!(self, Instruction::Subscript);
                }
                ExprContext::Store => {
                    // CPython uses STORE_SLICE
                    emit!(self, Instruction::BuildSlice { step: n == 3 });
                    emit!(self, Instruction::StoreSubscript);
                }
                _ => unreachable!(),
            }
        } else {
            // VISIT(c, expr, e->v.Subscript.slice)
            self.compile_expression(slice)?;

            // Emit appropriate instruction based on context
            match ctx {
                ExprContext::Load => emit!(self, Instruction::Subscript),
                ExprContext::Store => emit!(self, Instruction::StoreSubscript),
                ExprContext::Del => emit!(self, Instruction::DeleteSubscript),
                ExprContext::Invalid => {
                    return Err(self.error(CodegenErrorType::SyntaxError(
                        "Invalid expression context".to_owned(),
                    )));
                }
            }
        }

        Ok(())
    }

    /// Helper function for compiling tuples/lists/sets with starred expressions
    ///
    /// Parameters:
    /// - elts: The elements to compile
    /// - pushed: Number of items already on the stack
    /// - collection_type: What type of collection to build (tuple, list, set)
    ///
    // = starunpack_helper in compile.c
    fn starunpack_helper(
        &mut self,
        elts: &[Expr],
        pushed: u32,
        collection_type: CollectionType,
    ) -> CompileResult<()> {
        // Use RustPython's existing approach with BuildXFromTuples
        let (size, unpack) = self.gather_elements(pushed, elts)?;

        if unpack {
            // Has starred elements
            match collection_type {
                CollectionType::Tuple => {
                    if size > 1 || pushed > 0 {
                        emit!(self, Instruction::BuildTupleFromTuples { size });
                    }
                    // If size == 1 and pushed == 0, the single tuple is already on the stack
                }
                CollectionType::List => {
                    emit!(self, Instruction::BuildListFromTuples { size });
                }
                CollectionType::Set => {
                    emit!(self, Instruction::BuildSetFromTuples { size });
                }
            }
        } else {
            // No starred elements
            match collection_type {
                CollectionType::Tuple => {
                    emit!(self, Instruction::BuildTuple { size });
                }
                CollectionType::List => {
                    emit!(self, Instruction::BuildList { size });
                }
                CollectionType::Set => {
                    emit!(self, Instruction::BuildSet { size });
                }
            }
        }

        Ok(())
    }

    fn error(&mut self, error: CodegenErrorType) -> CodegenError {
        self.error_ranged(error, self.current_source_range)
    }

    fn error_ranged(&mut self, error: CodegenErrorType, range: TextRange) -> CodegenError {
        let location = self
            .source_file
            .to_source_code()
            .source_location(range.start(), PositionEncoding::Utf8);
        CodegenError {
            error,
            location: Some(location),
            source_path: self.source_file.name().to_owned(),
        }
    }

    /// Get the SymbolTable for the current scope.
    fn current_symbol_table(&self) -> &SymbolTable {
        if self.symbol_table_stack.is_empty() {
            panic!("symbol_table_stack is empty! This is a compiler bug.");
        }
        let index = self.symbol_table_stack.len() - 1;
        &self.symbol_table_stack[index]
    }

    /// Get the index of a free variable.
    fn get_free_var_index(&mut self, name: &str) -> CompileResult<u32> {
        let info = self.code_stack.last_mut().unwrap();
        let idx = info
            .metadata
            .freevars
            .get_index_of(name)
            .unwrap_or_else(|| info.metadata.freevars.insert_full(name.to_owned()).0);
        Ok((idx + info.metadata.cellvars.len()).to_u32())
    }

    /// Get the index of a cell variable.
    fn get_cell_var_index(&mut self, name: &str) -> CompileResult<u32> {
        let info = self.code_stack.last_mut().unwrap();
        let idx = info
            .metadata
            .cellvars
            .get_index_of(name)
            .unwrap_or_else(|| info.metadata.cellvars.insert_full(name.to_owned()).0);
        Ok(idx.to_u32())
    }

    /// Get the index of a local variable.
    fn get_local_var_index(&mut self, name: &str) -> CompileResult<u32> {
        let info = self.code_stack.last_mut().unwrap();
        let idx = info
            .metadata
            .varnames
            .get_index_of(name)
            .unwrap_or_else(|| info.metadata.varnames.insert_full(name.to_owned()).0);
        Ok(idx.to_u32())
    }

    /// Get the index of a global name.
    fn get_global_name_index(&mut self, name: &str) -> u32 {
        let info = self.code_stack.last_mut().unwrap();
        let idx = info
            .metadata
            .names
            .get_index_of(name)
            .unwrap_or_else(|| info.metadata.names.insert_full(name.to_owned()).0);
        idx.to_u32()
    }

    /// Push the next symbol table on to the stack
    fn push_symbol_table(&mut self) -> &SymbolTable {
        // Look up the next table contained in the scope of the current table
        let current_table = self
            .symbol_table_stack
            .last_mut()
            .expect("no current symbol table");

        if current_table.sub_tables.is_empty() {
            panic!(
                "push_symbol_table: no sub_tables available in {} (type: {:?})",
                current_table.name, current_table.typ
            );
        }

        let table = current_table.sub_tables.remove(0);

        // Push the next table onto the stack
        let last_idx = self.symbol_table_stack.len();
        self.symbol_table_stack.push(table);
        &self.symbol_table_stack[last_idx]
    }

    /// Pop the current symbol table off the stack
    fn pop_symbol_table(&mut self) -> SymbolTable {
        self.symbol_table_stack.pop().expect("compiler bug")
    }

    /// Enter a new scope
    // = compiler_enter_scope
    fn enter_scope(
        &mut self,
        name: &str,
        scope_type: CompilerScope,
        key: usize, // In RustPython, we use the index in symbol_table_stack as key
        lineno: u32,
    ) -> CompileResult<()> {
        // Create location
        let location = SourceLocation {
            line: OneIndexed::new(lineno as usize).unwrap_or(OneIndexed::MIN),
            character_offset: OneIndexed::MIN,
        };

        // Allocate a new compiler unit

        // In Rust, we'll create the structure directly
        let source_path = self.source_file.name().to_owned();

        // Lookup symbol table entry using key (_PySymtable_Lookup)
        let ste = if key < self.symbol_table_stack.len() {
            &self.symbol_table_stack[key]
        } else {
            return Err(self.error(CodegenErrorType::SyntaxError(
                "unknown symbol table entry".to_owned(),
            )));
        };

        // Use varnames from symbol table (already collected in definition order)
        let varname_cache: IndexSet<String> = ste.varnames.iter().cloned().collect();

        // Build cellvars using dictbytype (CELL scope, sorted)
        let mut cellvar_cache = IndexSet::default();
        let mut cell_names: Vec<_> = ste
            .symbols
            .iter()
            .filter(|(_, s)| s.scope == SymbolScope::Cell)
            .map(|(name, _)| name.clone())
            .collect();
        cell_names.sort();
        for name in cell_names {
            cellvar_cache.insert(name);
        }

        // Handle implicit __class__ cell if needed
        if ste.needs_class_closure {
            // Cook up an implicit __class__ cell
            debug_assert_eq!(scope_type, CompilerScope::Class);
            cellvar_cache.insert("__class__".to_string());
        }

        // Handle implicit __classdict__ cell if needed
        if ste.needs_classdict {
            // Cook up an implicit __classdict__ cell
            debug_assert_eq!(scope_type, CompilerScope::Class);
            cellvar_cache.insert("__classdict__".to_string());
        }

        // Build freevars using dictbytype (FREE scope, offset by cellvars size)
        let mut freevar_cache = IndexSet::default();
        let mut free_names: Vec<_> = ste
            .symbols
            .iter()
            .filter(|(_, s)| {
                s.scope == SymbolScope::Free || s.flags.contains(SymbolFlags::FREE_CLASS)
            })
            .map(|(name, _)| name.clone())
            .collect();
        free_names.sort();
        for name in free_names {
            freevar_cache.insert(name);
        }

        // Initialize u_metadata fields
        let (flags, posonlyarg_count, arg_count, kwonlyarg_count) = match scope_type {
            CompilerScope::Module => (bytecode::CodeFlags::empty(), 0, 0, 0),
            CompilerScope::Class => (bytecode::CodeFlags::empty(), 0, 0, 0),
            CompilerScope::Function | CompilerScope::AsyncFunction | CompilerScope::Lambda => (
                bytecode::CodeFlags::NEW_LOCALS | bytecode::CodeFlags::IS_OPTIMIZED,
                0, // Will be set later in enter_function
                0, // Will be set later in enter_function
                0, // Will be set later in enter_function
            ),
            CompilerScope::Comprehension => (
                bytecode::CodeFlags::NEW_LOCALS | bytecode::CodeFlags::IS_OPTIMIZED,
                0,
                1, // comprehensions take one argument (.0)
                0,
            ),
            CompilerScope::TypeParams => (
                bytecode::CodeFlags::NEW_LOCALS | bytecode::CodeFlags::IS_OPTIMIZED,
                0,
                0,
                0,
            ),
        };

        // Get private name from parent scope
        let private = if !self.code_stack.is_empty() {
            self.code_stack.last().unwrap().private.clone()
        } else {
            None
        };

        // Create the new compilation unit
        let code_info = ir::CodeInfo {
            flags,
            source_path: source_path.clone(),
            private,
            blocks: vec![ir::Block::default()],
            current_block: BlockIdx(0),
            metadata: ir::CodeUnitMetadata {
                name: name.to_owned(),
                qualname: None, // Will be set below
                consts: IndexSet::default(),
                names: IndexSet::default(),
                varnames: varname_cache,
                cellvars: cellvar_cache,
                freevars: freevar_cache,
                fast_hidden: IndexMap::default(),
                argcount: arg_count,
                posonlyargcount: posonlyarg_count,
                kwonlyargcount: kwonlyarg_count,
                firstlineno: OneIndexed::new(lineno as usize).unwrap_or(OneIndexed::MIN),
            },
            static_attributes: if scope_type == CompilerScope::Class {
                Some(IndexSet::default())
            } else {
                None
            },
            in_inlined_comp: false,
            fblock: Vec::with_capacity(MAXBLOCKS),
            symbol_table_index: key,
        };

        // Push the old compiler unit on the stack (like PyCapsule)
        // This happens before setting qualname
        self.code_stack.push(code_info);

        // Set qualname after pushing (uses compiler_set_qualname logic)
        if scope_type != CompilerScope::Module {
            self.set_qualname();
        }

        // Emit RESUME instruction
        let _resume_loc = if scope_type == CompilerScope::Module {
            // Module scope starts with lineno 0
            SourceLocation {
                line: OneIndexed::MIN,
                character_offset: OneIndexed::MIN,
            }
        } else {
            location
        };

        // Set the source range for the RESUME instruction
        // For now, just use an empty range at the beginning
        self.current_source_range = TextRange::default();
        emit!(
            self,
            Instruction::Resume {
                arg: bytecode::ResumeType::AtFuncStart as u32
            }
        );

        if scope_type == CompilerScope::Module {
            // This would be loc.lineno = -1 in CPython
            // We handle this differently in RustPython
        }

        Ok(())
    }

    fn push_output(
        &mut self,
        flags: bytecode::CodeFlags,
        posonlyarg_count: u32,
        arg_count: u32,
        kwonlyarg_count: u32,
        obj_name: String,
    ) {
        // First push the symbol table
        let table = self.push_symbol_table();
        let scope_type = table.typ;

        // The key is the current position in the symbol table stack
        let key = self.symbol_table_stack.len() - 1;

        // Get the line number
        let lineno = self.get_source_line_number().get();

        // Call enter_scope which does most of the work
        if let Err(e) = self.enter_scope(&obj_name, scope_type, key, lineno.to_u32()) {
            // In the current implementation, push_output doesn't return an error,
            // so we panic here. This maintains the same behavior.
            panic!("enter_scope failed: {e:?}");
        }

        // Override the values that push_output sets explicitly
        // enter_scope sets default values based on scope_type, but push_output
        // allows callers to specify exact values
        if let Some(info) = self.code_stack.last_mut() {
            info.flags = flags;
            info.metadata.argcount = arg_count;
            info.metadata.posonlyargcount = posonlyarg_count;
            info.metadata.kwonlyargcount = kwonlyarg_count;
        }
    }

    // compiler_exit_scope
    fn exit_scope(&mut self) -> CodeObject {
        let _table = self.pop_symbol_table();

        // Various scopes can have sub_tables:
        // - TypeParams scope can have sub_tables (the function body's symbol table)
        // - Module scope can have sub_tables (for TypeAlias scopes, nested functions, classes)
        // - Function scope can have sub_tables (for nested functions, classes)
        // - Class scope can have sub_tables (for nested classes, methods)

        let pop = self.code_stack.pop();
        let stack_top = compiler_unwrap_option(self, pop);
        // No parent scope stack to maintain
        unwrap_internal(self, stack_top.finalize_code(self.opts.optimize))
    }

    /// Push a new fblock
    // = compiler_push_fblock
    fn push_fblock(
        &mut self,
        fb_type: FBlockType,
        fb_block: BlockIdx,
        fb_exit: BlockIdx,
    ) -> CompileResult<()> {
        let code = self.current_code_info();
        if code.fblock.len() >= MAXBLOCKS {
            return Err(self.error(CodegenErrorType::SyntaxError(
                "too many statically nested blocks".to_owned(),
            )));
        }
        code.fblock.push(FBlockInfo {
            fb_type,
            fb_block,
            fb_exit,
        });
        Ok(())
    }

    /// Pop an fblock
    // = compiler_pop_fblock
    fn pop_fblock(&mut self, _expected_type: FBlockType) -> FBlockInfo {
        let code = self.current_code_info();
        // TODO: Add assertion to check expected type matches
        // assert!(matches!(fblock.fb_type, expected_type));
        code.fblock.pop().expect("fblock stack underflow")
    }

    // could take impl Into<Cow<str>>, but everything is borrowed from ast structs; we never
    // actually have a `String` to pass
    fn name(&mut self, name: &str) -> bytecode::NameIdx {
        self._name_inner(name, |i| &mut i.metadata.names)
    }
    fn varname(&mut self, name: &str) -> CompileResult<bytecode::NameIdx> {
        if Self::is_forbidden_arg_name(name) {
            return Err(self.error(CodegenErrorType::SyntaxError(format!(
                "cannot assign to {name}",
            ))));
        }
        Ok(self._name_inner(name, |i| &mut i.metadata.varnames))
    }
    fn _name_inner(
        &mut self,
        name: &str,
        cache: impl FnOnce(&mut ir::CodeInfo) -> &mut IndexSet<String>,
    ) -> bytecode::NameIdx {
        let name = self.mangle(name);
        let cache = cache(self.current_code_info());
        cache
            .get_index_of(name.as_ref())
            .unwrap_or_else(|| cache.insert_full(name.into_owned()).0)
            .to_u32()
    }

    /// Set the qualified name for the current code object
    // = compiler_set_qualname
    fn set_qualname(&mut self) -> String {
        let qualname = self.make_qualname();
        self.current_code_info().metadata.qualname = Some(qualname.clone());
        qualname
    }
    fn make_qualname(&mut self) -> String {
        let stack_size = self.code_stack.len();
        assert!(stack_size >= 1);

        let current_obj_name = self.current_code_info().metadata.name.clone();

        // If we're at the module level (stack_size == 1), qualname is just the name
        if stack_size <= 1 {
            return current_obj_name;
        }

        // Check parent scope
        let mut parent_idx = stack_size - 2;
        let mut parent = &self.code_stack[parent_idx];

        // If parent is TypeParams scope, look at grandparent
        // Check if parent is a type params scope by name pattern
        if parent.metadata.name.starts_with("<generic parameters of ") {
            if stack_size == 2 {
                // If we're immediately within the module, qualname is just the name
                return current_obj_name;
            }
            // Use grandparent
            parent_idx = stack_size - 3;
            parent = &self.code_stack[parent_idx];
        }

        // Check if this is a global class/function
        let mut force_global = false;
        if stack_size > self.symbol_table_stack.len() {
            // We might be in a situation where symbol table isn't pushed yet
            // In this case, check the parent symbol table
            if let Some(parent_table) = self.symbol_table_stack.last()
                && let Some(symbol) = parent_table.lookup(&current_obj_name)
                && symbol.scope == SymbolScope::GlobalExplicit
            {
                force_global = true;
            }
        } else if let Some(_current_table) = self.symbol_table_stack.last() {
            // Mangle the name if necessary (for private names in classes)
            let mangled_name = self.mangle(&current_obj_name);

            // Look up in parent symbol table to check scope
            if self.symbol_table_stack.len() >= 2 {
                let parent_table = &self.symbol_table_stack[self.symbol_table_stack.len() - 2];
                if let Some(symbol) = parent_table.lookup(&mangled_name)
                    && symbol.scope == SymbolScope::GlobalExplicit
                {
                    force_global = true;
                }
            }
        }

        // Build the qualified name
        if force_global {
            // For global symbols, qualname is just the name
            current_obj_name
        } else {
            // Check parent scope type
            let parent_obj_name = &parent.metadata.name;

            // Determine if parent is a function-like scope
            let is_function_parent = parent.flags.contains(bytecode::CodeFlags::IS_OPTIMIZED)
                && !parent_obj_name.starts_with("<") // Not a special scope like <lambda>, <listcomp>, etc.
                && parent_obj_name != "<module>"; // Not the module scope

            if is_function_parent {
                // For functions, append .<locals> to parent qualname
                // Use parent's qualname if available, otherwise use parent_obj_name
                let parent_qualname = parent.metadata.qualname.as_ref().unwrap_or(parent_obj_name);
                format!("{parent_qualname}.<locals>.{current_obj_name}")
            } else {
                // For classes and other scopes, use parent's qualname directly
                // Use parent's qualname if available, otherwise use parent_obj_name