ghostscope-compiler 0.1.5

//! eBPF LLVM context and core infrastructure
//!
//! This module provides the main code generation context and basic LLVM
//! infrastructure for eBPF program generation.

use super::maps::MapManager;
use crate::script::{VarType, VariableContext};
use ghostscope_dwarf::DwarfAnalyzer;
use inkwell::basic_block::BasicBlock;
use inkwell::builder::Builder;
use inkwell::context::Context;
use inkwell::debug_info::DebugInfoBuilder;
use inkwell::module::Module;
use inkwell::targets::{Target, TargetTriple};
use inkwell::values::{FunctionValue, IntValue, PointerValue};
use inkwell::AddressSpace;
use inkwell::OptimizationLevel;
use std::collections::HashMap;
use thiserror::Error;
use tracing::info;

/// Compile-time context containing PC address and module information for DWARF queries
#[derive(Debug, Clone)]
pub struct CompileTimeContext {
    pub pc_address: u64,
    pub module_path: String,
}

#[derive(Debug, Clone)]
pub struct BacktraceTailCallProgram {
    pub step_program_name: String,
}

#[derive(Debug, Clone)]
pub(crate) struct PendingBacktraceTailCall {
    pub step_program_name: String,
    pub depth: u8,
    pub instruction_size: usize,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct BacktraceModuleRowRangeEntry {
    pub cookie: u64,
    pub range: ghostscope_protocol::BacktraceModuleRowRange,
}

#[derive(Error, Debug)]
pub enum CodeGenError {
    #[error("LLVM compilation error: {0}")]
    LLVMError(String),
    #[error("Unsupported evaluation result: {0}")]
    UnsupportedEvaluation(String),
    #[error("Register mapping error: {0}")]
    RegisterMappingError(String),
    #[error("Memory access error: {0}")]
    MemoryAccessError(String),
    #[error("Builder error: {0}")]
    Builder(String),

    // === Variable lookup and availability errors ===
    #[error("Variable not found: {0}")]
    VariableNotFound(String),
    #[error("Variable not in scope: {0}")]
    VariableNotInScope(String),
    #[error("Variable unavailable: {0}")]
    VariableUnavailable(String),
    #[error("Type error: {0}")]
    TypeError(String),
    #[error("Not implemented: {0}")]
    NotImplemented(String),
    #[error("DWARF expression error: {0}")]
    DwarfError(String),
    #[error("Type size not available for variable: {0}")]
    TypeSizeNotAvailable(String),
}

pub type Result<T> = std::result::Result<T, CodeGenError>;

/// Runtime address produced by lowering a DWARF `PlannedAddress`.
///
/// `value` is the address to use in generated eBPF code. `offsets_found` is an
/// i1 guard that is false when a link-time address could not be rebased through
/// `proc_module_offsets`. Runtime-derived addresses do not need that lookup, so
/// their guard is always true.
#[derive(Debug, Clone, Copy)]
pub(crate) struct RuntimeAddress<'ctx> {
    pub value: IntValue<'ctx>,
    pub offsets_found: IntValue<'ctx>,
}

impl<'ctx> RuntimeAddress<'ctx> {
    pub(crate) fn available(value: IntValue<'ctx>, context: &'ctx Context) -> Self {
        Self {
            value,
            offsets_found: context.bool_type().const_int(1, false),
        }
    }

    pub(crate) fn with_offsets_found(value: IntValue<'ctx>, offsets_found: IntValue<'ctx>) -> Self {
        Self {
            value,
            offsets_found,
        }
    }

    pub(crate) fn with_value(self, value: IntValue<'ctx>) -> Self {
        Self { value, ..self }
    }
}

/// eBPF LLVM code generation context
pub struct EbpfContext<'ctx, 'dw> {
    pub context: &'ctx Context,
    pub module: Module<'ctx>,
    pub builder: Builder<'ctx>,

    // eBPF-specific function declarations
    pub trace_printk_fn: FunctionValue<'ctx>,

    // Map manager for eBPF maps
    pub map_manager: MapManager<'ctx>,

    // Debug infrastructure
    pub di_builder: DebugInfoBuilder<'ctx>,
    pub compile_unit: inkwell::debug_info::DICompileUnit<'ctx>,

    // === Complete Variable Management System ===
    pub variables: HashMap<String, PointerValue<'ctx>>, // Variable name -> LLVM pointer
    pub var_types: HashMap<String, VarType>,            // Variable name -> type
    pub optimized_out_vars: HashMap<String, bool>,      // Optimized out variables
    pub var_pc_addresses: HashMap<String, u64>,         // Variable -> PC address
    pub variable_context: Option<VariableContext>,      // Scope validation context
    pub(super) process_analyzer: Option<&'dw DwarfAnalyzer>, // Multi-module DWARF analyzer
    pub current_trace_id: Option<u32>,                  // Current trace_id being compiled
    pub current_compile_time_context: Option<CompileTimeContext>, // PC address and module for DWARF queries

    // === New instruction-based compilation system ===
    pub trace_context: ghostscope_protocol::TraceContext, // Trace context for optimized transmission

    // Per-invocation stack key for proc_module_offsets lookups (allocated in entry block)
    // Backed by `[4 x i32]`, so consumers may only assume i32 alignment.
    pub pm_key_alloca: Option<inkwell::values::PointerValue<'ctx>>,
    // Per-invocation 8-byte scratch for thread-pointer/TLS helper reads,
    // allocated lazily only for programs that actually lower TLS variables.
    pub(super) tls_scratch_alloca: Option<inkwell::values::PointerValue<'ctx>>,
    // Per-invocation event accumulation offset (u32) stored on stack (entry block)
    pub event_offset_alloca: Option<inkwell::values::PointerValue<'ctx>>,
    // Compile-time upper bound for bytes that may already be reserved in the current trace event.
    // This is maintained across structured control flow so later instructions can budget against
    // the worst-case path without double-counting sibling branches.
    pub compile_time_event_bytes_upper_bound: usize,
    // Compilation options (includes eBPF map configuration)
    pub compile_options: crate::CompileOptions,

    // === Control-flow expression error capture (soft abort) ===
    pub condition_context_active: bool,

    // === DWARF alias variables (script-level symbolic references) ===
    // These variables do not store pointer values; instead they remember the RHS
    // expression and are resolved to runtime addresses at use sites.
    pub alias_vars: HashMap<String, crate::script::Expr>,

    // === Script string variables (store literal bytes for content printing) ===
    // When a variable is bound from a string literal (or copied from another string var),
    // we keep its bytes (including optional NUL) here for content printing via ImmediateBytes.
    pub string_vars: HashMap<String, Vec<u8>>,

    // === DWARF compact unwind rows for bt ===
    pub backtrace_unwind_rows: Vec<ghostscope_protocol::BacktraceUnwindRow>,
    pub(crate) backtrace_module_row_ranges: Vec<BacktraceModuleRowRangeEntry>,
    pub(crate) backtrace_tail_call_slots: u8,
    pub(crate) next_backtrace_tail_call_slot: u8,
    pub(crate) pending_backtrace_tail_call: Option<PendingBacktraceTailCall>,
    pub(crate) backtrace_tail_enabled_alloca: Option<inkwell::values::PointerValue<'ctx>>,
    pub(crate) backtrace_tail_last_slot_alloca: Option<inkwell::values::PointerValue<'ctx>>,

    // === Lexical scoping for immutable variables ===
    // Each scope frame records names declared in that scope.
    pub scope_stack: Vec<std::collections::HashSet<String>>,
}

impl<'ctx, 'dw> EbpfContext<'ctx, 'dw> {
    pub(crate) fn backtrace_unwind_row_map_entries(&self) -> u64 {
        (self.compile_options.backtrace_unwind_rows_max_entries as u64)
            .max(self.backtrace_unwind_rows.len() as u64)
            .max(1)
    }

    /// Create a new eBPF code generation context
    pub fn new(
        context: &'ctx Context,
        module_name: &str,
        trace_id: Option<u32>,
        compile_options: &crate::CompileOptions,
    ) -> Result<Self> {
        let module = context.create_module(module_name);
        let builder = context.create_builder();

        // Initialize standard BPF target
        Target::initialize_bpf(&Default::default());

        // Create BPF target triple
        let triple = TargetTriple::create("bpf-pc-linux");

        // Get target and create target machine
        let target = Target::from_triple(&triple).map_err(|e| {
            CodeGenError::LLVMError(format!("Failed to get target from triple: {e}"))
        })?;
        let target_machine = target
            .create_target_machine(
                &triple,
                "generic",
                "+alu32",
                OptimizationLevel::Default,
                inkwell::targets::RelocMode::PIC,
                inkwell::targets::CodeModel::Small,
            )
            .ok_or_else(|| {
                CodeGenError::LLVMError("Failed to create target machine".to_string())
            })?;

        // Set module data layout and triple
        let data_layout = target_machine.get_target_data().get_data_layout();
        module.set_data_layout(&data_layout);
        module.set_triple(&triple);

        // Initialize debug info
        let (di_builder, compile_unit) = module.create_debug_info_builder(
            true,                                         // allow_unresolved
            inkwell::debug_info::DWARFSourceLanguage::C,  // language
            "ghostscope_generated.c",                     // filename
            ".",                                          // directory
            "ghostscope-compiler",                        // producer
            false,                                        // is_optimized
            "",                                           // flags
            1,                                            // runtime_version
            "",                                           // split_name
            inkwell::debug_info::DWARFEmissionKind::Full, // kind
            0,                                            // dwo_id
            false,                                        // split_debug_inlining
            false,                                        // debug_info_for_profiling
            "",                                           // sysroot
            "",                                           // sdk
        );

        let map_manager = MapManager::new(context);

        // Declare eBPF helper functions
        let trace_printk_fn = Self::declare_trace_printk(context, &module);

        Ok(Self {
            context,
            module,
            builder,
            trace_printk_fn,
            map_manager,
            di_builder,
            compile_unit,

            // Initialize variable management system
            variables: HashMap::new(),
            var_types: HashMap::new(),
            optimized_out_vars: HashMap::new(),
            var_pc_addresses: HashMap::new(),
            variable_context: None,
            process_analyzer: None,
            current_trace_id: trace_id,
            current_compile_time_context: None,

            // Initialize new instruction-based compilation system
            trace_context: ghostscope_protocol::TraceContext::new(),
            pm_key_alloca: None,
            tls_scratch_alloca: None,
            event_offset_alloca: None,
            compile_time_event_bytes_upper_bound: 0,
            compile_options: compile_options.clone(),

            // Control-flow expression context
            condition_context_active: false,

            // Alias variables
            alias_vars: HashMap::new(),
            // String variables
            string_vars: HashMap::new(),
            // Backtrace compact unwind rows
            backtrace_unwind_rows: Vec::new(),
            backtrace_module_row_ranges: Vec::new(),
            backtrace_tail_call_slots: 1,
            next_backtrace_tail_call_slot: 0,
            pending_backtrace_tail_call: None,
            backtrace_tail_enabled_alloca: None,
            backtrace_tail_last_slot_alloca: None,

            // Scopes
            scope_stack: Vec::new(),
        })
    }

    /// Enter a new lexical scope
    pub fn enter_scope(&mut self) {
        self.scope_stack.push(std::collections::HashSet::new());
    }

    /// Exit current lexical scope and drop all names declared within
    pub fn exit_scope(&mut self) {
        if let Some(names) = self.scope_stack.pop() {
            for name in names {
                self.variables.remove(&name);
                self.var_types.remove(&name);
                self.alias_vars.remove(&name);
                self.string_vars.remove(&name);
                self.optimized_out_vars.remove(&name);
                self.var_pc_addresses.remove(&name);
            }
        }
    }

    /// Check if a name exists in any active scope
    pub fn is_name_in_any_scope(&self, name: &str) -> bool {
        self.scope_stack.iter().any(|s| s.contains(name))
    }

    /// Check if a name exists in current (top) scope
    pub fn is_name_in_current_scope(&self, name: &str) -> bool {
        match self.scope_stack.last() {
            Some(top) => top.contains(name),
            None => false,
        }
    }

    /// Declare a name in the current scope. Disallow same-scope redeclaration and shadowing.
    pub fn declare_name_in_current_scope(&mut self, name: &str) -> Result<()> {
        if self.scope_stack.is_empty() {
            // Initialize a root scope if not present
            self.enter_scope();
        }
        if self.is_name_in_current_scope(name) {
            return Err(CodeGenError::TypeError(format!(
                "Redeclaration in the same scope is not allowed: '{name}'"
            )));
        }
        if self.is_name_in_any_scope(name) {
            return Err(CodeGenError::TypeError(format!(
                "Shadowing is not allowed for immutable variables: '{name}'"
            )));
        }
        if let Some(top) = self.scope_stack.last_mut() {
            top.insert(name.to_string());
        }
        Ok(())
    }

    /// Create a new code generator with DWARF analyzer support
    pub fn new_with_process_analyzer(
        context: &'ctx Context,
        module_name: &str,
        process_analyzer: Option<&'dw DwarfAnalyzer>,
        trace_id: Option<u32>,
        compile_options: &crate::CompileOptions,
    ) -> Result<Self> {
        let mut codegen = Self::new(context, module_name, trace_id, compile_options)?;
        codegen.process_analyzer = process_analyzer;
        Ok(codegen)
    }

    /// Set compile-time context for DWARF queries
    pub fn set_compile_time_context(&mut self, pc_address: u64, module_path: String) {
        self.current_compile_time_context = Some(CompileTimeContext {
            pc_address,
            module_path,
        });
    }

    /// Get compile-time context for DWARF queries
    pub fn get_compile_time_context(&self) -> Result<&CompileTimeContext> {
        self.current_compile_time_context
            .as_ref()
            .ok_or_else(|| CodeGenError::DwarfError("No compile-time context set".to_string()))
    }

    /// Declare trace_printk eBPF helper function
    fn declare_trace_printk(context: &'ctx Context, module: &Module<'ctx>) -> FunctionValue<'ctx> {
        let i32_type = context.i32_type();
        let ptr_type = context.ptr_type(AddressSpace::default());
        let i64_type = context.i64_type();

        // int bpf_trace_printk(const char *fmt, u32 fmt_size, ...)
        let fn_type = i32_type.fn_type(&[ptr_type.into(), i64_type.into()], true);

        module.add_function("bpf_trace_printk", fn_type, None)
    }

    /// Create basic eBPF function with proper signature
    pub fn create_basic_ebpf_function(&mut self, function_name: &str) -> Result<()> {
        let i32_type = self.context.i32_type();
        let ptr_type = self.context.ptr_type(AddressSpace::default());

        // eBPF function signature: int function(struct pt_regs *ctx)
        let fn_type = i32_type.fn_type(&[ptr_type.into()], false);

        let function = self.module.add_function(function_name, fn_type, None);

        // Set section attribute for uprobe
        function.add_attribute(
            inkwell::attributes::AttributeLoc::Function,
            self.context.create_string_attribute("section", "uprobe"),
        );

        // Create basic block
        let basic_block = self.context.append_basic_block(function, "entry");
        self.builder.position_at_end(basic_block);

        info!("Created eBPF function: {}", function_name);
        Ok(())
    }

    /// Test helper: ensure proc_module_offsets map exists in the module
    #[cfg(test)]
    pub fn __test_ensure_proc_offsets_map(&mut self) -> Result<()> {
        self.map_manager
            .create_proc_module_offsets_map(
                &self.module,
                &self.di_builder,
                &self.compile_unit,
                "proc_module_offsets",
                self.compile_options.proc_module_offsets_max_entries,
            )
            .map_err(|e| {
                CodeGenError::LLVMError(format!(
                    "Failed to create proc_module_offsets map in test: {e}"
                ))
            })?;
        self.map_manager
            .create_pid_aliases_map(
                &self.module,
                &self.di_builder,
                &self.compile_unit,
                "pid_aliases",
                self.compile_options.proc_module_offsets_max_entries,
            )
            .map_err(|e| {
                CodeGenError::LLVMError(format!("Failed to create pid_aliases map in test: {e}"))
            })?;
        self.map_manager
            .create_proc_module_range_meta_map(
                &self.module,
                &self.di_builder,
                &self.compile_unit,
                "proc_module_range_meta",
                self.compile_options.proc_module_offsets_max_entries,
            )
            .map_err(|e| {
                CodeGenError::LLVMError(format!(
                    "Failed to create proc_module_range_meta map in test: {e}"
                ))
            })?;
        self.map_manager
            .create_proc_module_ranges_map(
                &self.module,
                &self.di_builder,
                &self.compile_unit,
                "proc_module_ranges",
                self.compile_options
                    .proc_module_offsets_max_entries
                    .saturating_mul(2)
                    .max(1),
            )
            .map_err(|e| {
                CodeGenError::LLVMError(format!(
                    "Failed to create proc_module_ranges map in test: {e}"
                ))
            })
    }

    /// Test helper: allocate per-invocation pm_key on the entry block like create_main_function
    #[cfg(test)]
    pub fn __test_alloc_pm_key(&mut self) -> Result<()> {
        let i32_type = self.context.i32_type();
        let key_arr_ty = i32_type.array_type(4);
        let key_alloca = self
            .builder
            .build_alloca(key_arr_ty, "pm_key")
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
        self.pm_key_alloca = Some(key_alloca);
        Ok(())
    }

    /// Get the LLVM module reference
    pub fn get_module(&self) -> &Module<'ctx> {
        &self.module
    }

    /// Get the string table after compilation
    pub fn get_trace_context(&self) -> ghostscope_protocol::TraceContext {
        self.trace_context.clone()
    }

    pub fn backtrace_tail_call_program(&self) -> Option<BacktraceTailCallProgram> {
        self.pending_backtrace_tail_call
            .as_ref()
            .map(|plan| BacktraceTailCallProgram {
                step_program_name: plan.step_program_name.clone(),
            })
    }

    pub(crate) fn current_insert_block(&self, op: &str) -> Result<BasicBlock<'ctx>> {
        self.builder
            .get_insert_block()
            .ok_or_else(|| CodeGenError::Builder(format!("{op} requires an active insert block")))
    }

    pub(crate) fn current_function(&self, op: &str) -> Result<FunctionValue<'ctx>> {
        self.current_insert_block(op)?
            .get_parent()
            .ok_or_else(|| CodeGenError::Builder(format!("{op} requires a parent function")))
    }

    /// Get pt_regs parameter from current function
    pub fn get_pt_regs_parameter(&self) -> Result<PointerValue<'ctx>> {
        let current_function = self.current_function("get pt_regs parameter")?;

        let pt_regs_param = current_function
            .get_first_param()
            .ok_or_else(|| CodeGenError::Builder("Function has no parameters".to_string()))?
            .into_pointer_value();

        Ok(pt_regs_param)
    }

    /// Compile a complete program with statements
    pub fn compile_program(
        &mut self,
        _program: &crate::script::Program,
        function_name: &str,
        trace_statements: &[crate::script::Statement],
        target_pid: Option<u32>,
        compile_time_pc: Option<u64>,
        module_path: Option<&str>,
    ) -> Result<(FunctionValue<'ctx>, ghostscope_protocol::TraceContext)> {
        info!(
            "Starting program compilation with function: {}",
            function_name
        );

        // Set the current trace_id and compile-time context for code generation
        self.current_compile_time_context =
            if let (Some(pc), Some(path)) = (compile_time_pc, module_path) {
                Some(CompileTimeContext {
                    pc_address: pc,
                    module_path: path.to_string(),
                })
            } else {
                None
            };
        self.prepare_backtrace_unwind_rows(trace_statements);

        // Create required maps - critical for eBPF loader
        // Create event output map based on compile options
        match self.compile_options.event_map_type {
            crate::EventMapType::RingBuf => {
                self.map_manager
                    .create_ringbuf_map(
                        &self.module,
                        &self.di_builder,
                        &self.compile_unit,
                        "ringbuf",
                        self.compile_options.ringbuf_size,
                    )
                    .map_err(|e| {
                        CodeGenError::LLVMError(format!("Failed to create ringbuf map: {e}"))
                    })?;
            }
            crate::EventMapType::PerfEventArray => {
                self.map_manager
                    .create_perf_event_array_map(
                        &self.module,
                        &self.di_builder,
                        &self.compile_unit,
                        "events",
                    )
                    .map_err(|e| {
                        CodeGenError::LLVMError(format!(
                            "Failed to create perf event array map: {e}"
                        ))
                    })?;
            }
        }

        self.map_manager
            .create_event_loss_counter_map(
                &self.module,
                &self.di_builder,
                &self.compile_unit,
                "event_loss_counters",
                1,
            )
            .map_err(|e| {
                CodeGenError::LLVMError(format!("Failed to create event_loss_counters map: {e}"))
            })?;

        // Create per-CPU accumulation maps for single-record event emission
        //  - event_accum_buffer: value size = max_trace_event_size bytes, entries = 1
        //  - event_accum_offset: value size = 4 bytes (u32), entries = 1
        self.map_manager
            .create_percpu_array_map(
                &self.module,
                &self.di_builder,
                &self.compile_unit,
                "event_accum_buffer",
                1,
                self.compile_options.max_trace_event_size as u64,
            )
            .map_err(|e| {
                CodeGenError::LLVMError(format!("Failed to create event_accum_buffer: {e}"))
            })?;

        // Create ASLR offsets map for (pid,module) → section offsets
        self.map_manager
            .create_proc_module_offsets_map(
                &self.module,
                &self.di_builder,
                &self.compile_unit,
                "proc_module_offsets",
                self.compile_options.proc_module_offsets_max_entries,
            )
            .map_err(|e| {
                CodeGenError::LLVMError(format!("Failed to create proc_module_offsets map: {e}"))
            })?;

        self.map_manager
            .create_pid_aliases_map(
                &self.module,
                &self.di_builder,
                &self.compile_unit,
                "pid_aliases",
                self.compile_options.proc_module_offsets_max_entries,
            )
            .map_err(|e| {
                CodeGenError::LLVMError(format!("Failed to create pid_aliases map: {e}"))
            })?;

        if !self.backtrace_unwind_rows.is_empty() {
            let share_backtrace_maps = !self.backtrace_module_row_ranges.is_empty();
            if share_backtrace_maps {
                self.map_manager.mark_pinned_map("bt_unwind_rows");
                self.map_manager.mark_pinned_map("bt_module_row_ranges");
            }
            self.map_manager
                .create_array_map(
                    &self.module,
                    &self.di_builder,
                    &self.compile_unit,
                    "bt_unwind_rows",
                    self.backtrace_unwind_row_map_entries(),
                    crate::BACKTRACE_UNWIND_ROW_SIZE as u64,
                )
                .map_err(|e| {
                    CodeGenError::LLVMError(format!("Failed to create bt_unwind_rows map: {e}"))
                })?;
            if share_backtrace_maps {
                self.map_manager
                    .create_proc_module_range_meta_map(
                        &self.module,
                        &self.di_builder,
                        &self.compile_unit,
                        "proc_module_range_meta",
                        self.compile_options.proc_module_offsets_max_entries,
                    )
                    .map_err(|e| {
                        CodeGenError::LLVMError(format!(
                            "Failed to create proc_module_range_meta map: {e}"
                        ))
                    })?;
                self.map_manager
                    .create_proc_module_ranges_map(
                        &self.module,
                        &self.di_builder,
                        &self.compile_unit,
                        "proc_module_ranges",
                        self.compile_options
                            .proc_module_offsets_max_entries
                            .saturating_mul(2)
                            .max(1),
                    )
                    .map_err(|e| {
                        CodeGenError::LLVMError(format!(
                            "Failed to create proc_module_ranges map: {e}"
                        ))
                    })?;
                self.map_manager
                    .create_hash_map(
                        &self.module,
                        &self.di_builder,
                        &self.compile_unit,
                        "bt_module_row_ranges",
                        self.compile_options.proc_module_offsets_max_entries.max(1),
                        (
                            std::mem::size_of::<u64>() as u64,
                            ghostscope_protocol::BACKTRACE_MODULE_ROW_RANGE_SIZE as u64,
                        ),
                    )
                    .map_err(|e| {
                        CodeGenError::LLVMError(format!(
                            "Failed to create bt_module_row_ranges map: {e}"
                        ))
                    })?;
            }
            self.map_manager
                .create_percpu_array_map(
                    &self.module,
                    &self.di_builder,
                    &self.compile_unit,
                    "bt_state",
                    self.backtrace_tail_call_slots.max(1) as u64,
                    crate::BACKTRACE_TAIL_STATE_SIZE as u64,
                )
                .map_err(|e| {
                    CodeGenError::LLVMError(format!("Failed to create bt_state map: {e}"))
                })?;
            self.map_manager
                .create_program_array_map(
                    &self.module,
                    &self.di_builder,
                    &self.compile_unit,
                    "bt_prog_array",
                    1,
                )
                .map_err(|e| {
                    CodeGenError::LLVMError(format!("Failed to create bt_prog_array map: {e}"))
                })?;
        }

        // Variables are now queried on-demand when accessed in expressions
        // No need to pre-populate DWARF variables

        // Create main function
        let main_function = self.create_main_function(function_name)?;

        // Add PID filtering:
        // 1) explicit compile option override (namespace-aware)
        // 2) fallback to legacy host TGID filter from target_pid
        let pid_filter_spec = self
            .compile_options
            .pid_filter_spec
            .or_else(|| target_pid.map(|pid| crate::PidFilterSpec::HostTgid { filter_pid: pid }));
        if let Some(spec) = pid_filter_spec {
            self.add_pid_filter(spec)?;
        }

        // Use new staged transmission system for all statements
        let program = crate::script::ast::Program {
            statements: trace_statements.to_vec(),
        };

        // Collect variable types from DWARF analysis
        let variable_types = std::collections::HashMap::new(); // Empty for now, will be populated by codegen

        // Generate staged transmission code using new architecture
        let trace_context =
            self.compile_program_with_staged_transmission(&program, variable_types)?;
        info!(
            "Generated TraceContext with {} strings",
            trace_context.string_count()
        );

        // Return success
        let i32_type = self.context.i32_type();
        let return_value = i32_type.const_int(0, false);
        self.builder
            .build_return(Some(&return_value))
            .map_err(|e| CodeGenError::Builder(e.to_string()))?;

        info!(
            "Successfully compiled program with function: {} and TraceContext",
            function_name
        );
        Ok((main_function, trace_context))
    }

    /// Create the main eBPF function
    fn create_main_function(&mut self, function_name: &str) -> Result<FunctionValue<'ctx>> {
        let i32_type = self.context.i32_type();
        let ptr_type = self.context.ptr_type(AddressSpace::default());

        // Create function type: int function_name(void *ctx)
        let fn_type = i32_type.fn_type(&[ptr_type.into()], false);
        let function = self.module.add_function(function_name, fn_type, None);

        // CRITICAL: Set section name for eBPF loader to find the function
        function.set_section(Some("uprobe"));

        // Create basic block and position builder
        let basic_block = self.context.append_basic_block(function, "entry");
        self.builder.position_at_end(basic_block);

        // Allocate fixed-size per-invocation key buffer on the eBPF stack (entry block)
        // Layout: [ pid:u32, pad:u32, cookie_lo:u32, cookie_hi:u32 ] to match struct {u32; u64}
        // This remains an i32-aligned slot; probe-read scratch reuse must stay
        // limited to <=4-byte scalar loads.
        let key_arr_ty = i32_type.array_type(4);
        let key_alloca = self
            .builder
            .build_alloca(key_arr_ty, "pm_key")
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
        self.pm_key_alloca = Some(key_alloca);

        // Allocate per-invocation event_offset (u32) and initialize to 0
        let event_off_alloca = self
            .builder
            .build_alloca(i32_type, "event_offset")
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
        self.builder
            .build_store(event_off_alloca, i32_type.const_zero())
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
        self.event_offset_alloca = Some(event_off_alloca);

        info!("Created main function: {}", function_name);
        Ok(function)
    }

    pub(crate) fn create_tail_call_function(
        &mut self,
        function_name: &str,
    ) -> Result<FunctionValue<'ctx>> {
        let i32_type = self.context.i32_type();
        let ptr_type = self.context.ptr_type(AddressSpace::default());
        let fn_type = i32_type.fn_type(&[ptr_type.into()], false);
        let function = self.module.add_function(function_name, fn_type, None);
        function.set_section(Some("uprobe"));

        let basic_block = self.context.append_basic_block(function, "entry");
        self.builder.position_at_end(basic_block);

        let key_arr_ty = i32_type.array_type(4);
        let key_alloca = self
            .builder
            .build_alloca(key_arr_ty, "pm_key")
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
        self.pm_key_alloca = Some(key_alloca);

        info!("Created tail-call eBPF function: {}", function_name);
        Ok(function)
    }

    /// Add PID filtering logic to the current function.
    /// This generates LLVM IR to check PID and early-return if not matching.
    fn add_pid_filter(&mut self, spec: crate::PidFilterSpec) -> Result<()> {
        match spec {
            crate::PidFilterSpec::HostTgid { filter_pid } => self.add_host_pid_filter(filter_pid),
            crate::PidFilterSpec::NamespaceTgid { filter_pid, pid_ns } => {
                let (pid_ns_dev, pid_ns_inode) = pid_ns.helper_dev_inode().ok_or_else(|| {
                    CodeGenError::LLVMError(
                        "Namespace TGID filter requires pid namespace device id".to_string(),
                    )
                })?;
                self.add_namespace_pid_filter(filter_pid, pid_ns_dev, pid_ns_inode)
            }
        }
    }

    fn add_host_pid_filter(&mut self, filter_pid: u32) -> Result<()> {
        info!("Adding host TGID filter for filter PID: {}", filter_pid);

        // Get current function and entry block
        let current_fn = self.current_function("add host pid filter")?;

        // Create basic blocks for control flow
        let continue_block = self
            .context
            .append_basic_block(current_fn, "continue_execution");
        let early_return_block = self
            .context
            .append_basic_block(current_fn, "pid_mismatch_return");

        // Get current PID/TID using bpf_get_current_pid_tgid helper
        let pid_tgid_value = self.get_current_pid_tgid()?;

        // Extract TGID (high 32 bits) by right shifting 32 bits
        let shift_amount = self.context.i64_type().const_int(32, false);
        let current_tgid = self
            .builder
            .build_right_shift(pid_tgid_value, shift_amount, false, "current_tgid")
            .map_err(|e| CodeGenError::Builder(e.to_string()))?;

        // Convert filter_pid to i64 and compare
        let target_pid_value = self.context.i64_type().const_int(filter_pid as u64, false);
        let pid_matches = self
            .builder
            .build_int_compare(
                inkwell::IntPredicate::EQ,
                current_tgid,
                target_pid_value,
                "pid_matches",
            )
            .map_err(|e| CodeGenError::Builder(e.to_string()))?;

        // Conditional branch: if pid matches, continue; else early return
        self.builder
            .build_conditional_branch(pid_matches, continue_block, early_return_block)
            .map_err(|e| CodeGenError::Builder(e.to_string()))?;

        // Early return block - just return 0
        self.builder.position_at_end(early_return_block);
        self.builder
            .build_return(Some(&self.context.i32_type().const_int(0, false)))
            .map_err(|e| CodeGenError::Builder(e.to_string()))?;

        // Position at continue block for the rest of the function
        self.builder.position_at_end(continue_block);

        info!(
            "Host TGID filter added successfully for filter PID: {}",
            filter_pid
        );
        Ok(())
    }

    fn add_namespace_pid_filter(
        &mut self,
        filter_pid: u32,
        pid_ns_dev: u64,
        pid_ns_inode: u64,
    ) -> Result<()> {
        const BPF_FUNC_GET_NS_CURRENT_PID_TGID: u64 = 120;
        const BPF_PIDNS_INFO_SIZE: u64 = 8; // struct { u32 pid; u32 tgid; }

        info!(
            "Adding namespace TGID filter: filter_pid={} ns_dev={} ns_inode={}",
            filter_pid, pid_ns_dev, pid_ns_inode
        );

        let current_fn = self
            .builder
            .get_insert_block()
            .ok_or_else(|| CodeGenError::Builder("No current insert block".to_string()))?
            .get_parent()
            .ok_or_else(|| CodeGenError::Builder("No parent function".to_string()))?;

        let helper_ok_block = self
            .context
            .append_basic_block(current_fn, "pidns_helper_ok");
        let continue_block = self
            .context
            .append_basic_block(current_fn, "continue_execution");
        let early_return_block = self
            .context
            .append_basic_block(current_fn, "pid_mismatch_return");

        let i32_type = self.context.i32_type();
        let i64_type = self.context.i64_type();
        let ptr_type = self.context.ptr_type(AddressSpace::default());

        // Stack-allocate bpf_pidns_info-compatible storage: [pid:u32, tgid:u32].
        let pidns_info_ty = i32_type.array_type(2);
        let pidns_info_alloca = self
            .builder
            .build_alloca(pidns_info_ty, "pidns_info")
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
        self.builder
            .build_store(pidns_info_alloca, pidns_info_ty.const_zero())
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;

        let pidns_info_ptr = self
            .builder
            .build_bit_cast(pidns_info_alloca, ptr_type, "pidns_info_ptr")
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;

        let helper_args = [
            i64_type.const_int(pid_ns_dev, false).into(),
            i64_type.const_int(pid_ns_inode, false).into(),
            pidns_info_ptr,
            i64_type.const_int(BPF_PIDNS_INFO_SIZE, false).into(),
        ];
        let helper_ret = self.create_bpf_helper_call(
            BPF_FUNC_GET_NS_CURRENT_PID_TGID,
            &helper_args,
            i64_type.into(),
            "ns_pid_tgid_ret",
        )?;
        let helper_ret = match helper_ret {
            inkwell::values::BasicValueEnum::IntValue(v) => v,
            _ => {
                return Err(CodeGenError::LLVMError(
                    "bpf_get_ns_current_pid_tgid did not return integer".to_string(),
                ));
            }
        };

        let helper_ok = self
            .builder
            .build_int_compare(
                inkwell::IntPredicate::EQ,
                helper_ret,
                i64_type.const_zero(),
                "pidns_helper_ok",
            )
            .map_err(|e| CodeGenError::Builder(e.to_string()))?;
        self.builder
            .build_conditional_branch(helper_ok, helper_ok_block, early_return_block)
            .map_err(|e| CodeGenError::Builder(e.to_string()))?;

        self.builder.position_at_end(helper_ok_block);
        // SAFETY: pidns_info_alloca has the pid namespace helper result layout
        // [pid, tgid], so [0, 1] addresses the tgid field.
        let tgid_ptr = unsafe {
            self.builder.build_gep(
                pidns_info_ty,
                pidns_info_alloca,
                &[i32_type.const_zero(), i32_type.const_int(1, false)],
                "pidns_tgid_ptr",
            )
        }
        .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
        let ns_tgid = self
            .builder
            .build_load(i32_type, tgid_ptr, "ns_tgid")
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?
            .into_int_value();
        let ns_tgid_i64 = self
            .builder
            .build_int_z_extend(ns_tgid, i64_type, "ns_tgid_i64")
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
        let target_pid_value = i64_type.const_int(filter_pid as u64, false);
        let pid_matches = self
            .builder
            .build_int_compare(
                inkwell::IntPredicate::EQ,
                ns_tgid_i64,
                target_pid_value,
                "pid_matches_ns",
            )
            .map_err(|e| CodeGenError::Builder(e.to_string()))?;
        self.builder
            .build_conditional_branch(pid_matches, continue_block, early_return_block)
            .map_err(|e| CodeGenError::Builder(e.to_string()))?;

        self.builder.position_at_end(early_return_block);
        self.builder
            .build_return(Some(&self.context.i32_type().const_int(0, false)))
            .map_err(|e| CodeGenError::Builder(e.to_string()))?;

        self.builder.position_at_end(continue_block);
        info!(
            "Namespace TGID filter added successfully for filter PID: {}",
            filter_pid
        );
        Ok(())
    }

    /// Get or create a global i8 flag by name, initialized to 0
    pub fn get_or_create_flag_global(&mut self, name: &str) -> PointerValue<'ctx> {
        if let Some(g) = self.module.get_global(name) {
            return g.as_pointer_value();
        }
        let i8_type = self.context.i8_type();
        let global = self
            .module
            .add_global(i8_type, Some(AddressSpace::default()), name);
        global.set_initializer(&i8_type.const_zero());
        global.as_pointer_value()
    }

    /// Set a flag global to a constant u8 value at runtime
    pub fn store_flag_value(&mut self, name: &str, value: u8) -> Result<()> {
        let ptr = self.get_or_create_flag_global(name);
        self.builder
            .build_store(ptr, self.context.i8_type().const_int(value as u64, false))
            .map_err(|e| CodeGenError::LLVMError(format!("Failed to store flag {name}: {e}")))?;
        Ok(())
    }

    /// Mark that at least one variable succeeded (status==0)
    pub fn mark_any_success(&mut self) -> Result<()> {
        self.store_flag_value("_gs_any_success", 1)
    }

    /// Mark that at least one variable failed (status!=0)
    pub fn mark_any_fail(&mut self) -> Result<()> {
        self.store_flag_value("_gs_any_fail", 1)
    }

    /// Get or create global for condition error code (i8). Name: _gs_cond_error
    pub fn get_or_create_cond_error_global(&mut self) -> PointerValue<'ctx> {
        if let Some(g) = self.module.get_global("_gs_cond_error") {
            return g.as_pointer_value();
        }
        let i8_type = self.context.i8_type();
        let global =
            self.module
                .add_global(i8_type, Some(AddressSpace::default()), "_gs_cond_error");
        global.set_initializer(&i8_type.const_zero());
        global.as_pointer_value()
    }

    /// Reset condition error to 0 (only meaningful when condition_context_active=true)
    pub fn reset_condition_error(&mut self) -> Result<()> {
        let ptr = self.get_or_create_cond_error_global();
        self.builder
            .build_store(ptr, self.context.i8_type().const_zero())
            .map_err(|e| CodeGenError::LLVMError(format!("Failed to reset _gs_cond_error: {e}")))?;
        // Also reset error address
        let aptr = self.get_or_create_cond_error_addr_global();
        self.builder
            .build_store(aptr, self.context.i64_type().const_zero())
            .map_err(|e| {
                CodeGenError::LLVMError(format!("Failed to reset _gs_cond_error_addr: {e}"))
            })?;
        // Also reset flags
        let fptr = self.get_or_create_cond_error_flags_global();
        self.builder
            .build_store(fptr, self.context.i8_type().const_zero())
            .map_err(|e| {
                CodeGenError::LLVMError(format!("Failed to reset _gs_cond_error_flags: {e}"))
            })?;
        Ok(())
    }

    /// Get or create global for condition error address (i64). Name: _gs_cond_error_addr
    pub fn get_or_create_cond_error_addr_global(&mut self) -> PointerValue<'ctx> {
        if let Some(g) = self.module.get_global("_gs_cond_error_addr") {
            return g.as_pointer_value();
        }
        let i64_type = self.context.i64_type();
        let global = self.module.add_global(
            i64_type,
            Some(AddressSpace::default()),
            "_gs_cond_error_addr",
        );
        global.set_initializer(&i64_type.const_zero());
        global.as_pointer_value()
    }

    /// If in condition context, set error code when it's currently 0 (first error wins)
    pub fn set_condition_error_if_unset(&mut self, code: u8) -> Result<()> {
        if !self.condition_context_active {
            return Ok(());
        }
        let ptr = self.get_or_create_cond_error_global();
        let cur = self
            .builder
            .build_load(self.context.i8_type(), ptr, "cond_err_cur")
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?
            .into_int_value();
        let is_zero = self
            .builder
            .build_int_compare(
                inkwell::IntPredicate::EQ,
                cur,
                self.context.i8_type().const_zero(),
                "cond_err_is_zero",
            )
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
        let newv_bv: inkwell::values::BasicValueEnum =
            self.context.i8_type().const_int(code as u64, false).into();
        let sel = self
            .builder
            .build_select::<inkwell::values::BasicValueEnum, _>(
                is_zero,
                newv_bv,
                cur.into(),
                "cond_err_new",
            )
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
        self.builder
            .build_store(ptr, sel)
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
        Ok(())
    }

    /// Read the current condition error as i1 predicate: (error != 0)
    pub fn build_condition_error_predicate(&mut self) -> Result<inkwell::values::IntValue<'ctx>> {
        let ptr = self.get_or_create_cond_error_global();
        let cur = self
            .builder
            .build_load(self.context.i8_type(), ptr, "cond_err_cur")
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?
            .into_int_value();
        self.builder
            .build_int_compare(
                inkwell::IntPredicate::NE,
                cur,
                self.context.i8_type().const_zero(),
                "cond_err_nonzero",
            )
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))
    }

    /// Get or create global for condition error flags (i8). Name: _gs_cond_error_flags
    pub fn get_or_create_cond_error_flags_global(&mut self) -> PointerValue<'ctx> {
        if let Some(g) = self.module.get_global("_gs_cond_error_flags") {
            return g.as_pointer_value();
        }
        let i8_type = self.context.i8_type();
        let global = self.module.add_global(
            i8_type,
            Some(AddressSpace::default()),
            "_gs_cond_error_flags",
        );
        global.set_initializer(&i8_type.const_zero());
        global.as_pointer_value()
    }

    /// OR into condition error flags (BV must be i8)
    pub fn or_condition_error_flags(&mut self, flags: IntValue<'ctx>) -> Result<()> {
        if !self.condition_context_active {
            return Ok(());
        }
        let ptr = self.get_or_create_cond_error_flags_global();
        let cur = self
            .builder
            .build_load(self.context.i8_type(), ptr, "cond_err_flags_cur")
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?
            .into_int_value();
        let newv = self
            .builder
            .build_or(cur, flags, "cond_err_flags_or")
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
        self.builder
            .build_store(ptr, newv)
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
        Ok(())
    }

    /// If in condition context, record failing address (first win). addr must be i64
    pub fn set_condition_error_addr_if_unset(&mut self, addr: IntValue<'ctx>) -> Result<()> {
        if !self.condition_context_active {
            return Ok(());
        }
        let ptr = self.get_or_create_cond_error_addr_global();
        let cur = self
            .builder
            .build_load(self.context.i64_type(), ptr, "cond_err_addr_cur")
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?
            .into_int_value();
        let is_zero = self
            .builder
            .build_int_compare(
                inkwell::IntPredicate::EQ,
                cur,
                self.context.i64_type().const_zero(),
                "cond_err_addr_is_zero",
            )
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
        let sel = self
            .builder
            .build_select::<IntValue<'ctx>, _>(is_zero, addr, cur, "cond_err_addr_new")
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
        self.builder
            .build_store(ptr, sel)
            .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
        Ok(())
    }
}