cairo_lang_lowering/

objects.rs

//! Intermediate representation objects after lowering from semantic.
//!
//! This representation is SSA (static single-assignment): each variable is defined before usage and
//! assigned once. It is also normal form: each function argument is a variable, rather than a
//! compound expression.

use std::ops::{Deref, DerefMut};

use cairo_lang_debug::DebugWithDb;
use cairo_lang_defs::diagnostic_utils::StableLocation;
use cairo_lang_diagnostics::{DiagnosticNote, Diagnostics};
use cairo_lang_proc_macros::HeapSize;
use cairo_lang_semantic as semantic;
use cairo_lang_semantic::corelib::{concrete_destruct_trait, concrete_panic_destruct_trait};
use cairo_lang_semantic::expr::inference::InferenceError;
use cairo_lang_semantic::expr::inference::solver::Ambiguity;
use cairo_lang_semantic::items::constant::ConstValueId;
use cairo_lang_semantic::items::imp::ImplLookupContextId;
use cairo_lang_semantic::types::{TypeInfo, TypesSemantic};
use cairo_lang_semantic::{ConcreteEnumId, ConcreteVariant};
use cairo_lang_utils::Intern;
use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
use id_arena::{Arena, DefaultArenaBehavior, Id};

pub mod blocks;
pub use blocks::BlockId;
use salsa::Database;
use semantic::MatchArmSelector;

use self::blocks::Blocks;
use crate::diagnostic::LoweringDiagnostic;
use crate::fmt::LoweredFormatter;
use crate::ids::{FunctionId, LocationId, Signature};

/// The Location struct represents the source location of a lowered object. It is used to store the
/// most relevant source location for a lowering object.
#[derive(Clone, Debug, Eq, Hash, PartialEq, HeapSize, salsa::Update)]
pub struct Location<'db> {
    /// The stable location of the object.
    pub stable_location: StableLocation<'db>,
    /// Additional notes about the origin of the object, for example if the object was
    /// auto-generated by the compiler.
    /// New notes are appended to the end of the vector.
    pub notes: Vec<DiagnosticNote<'db>>,
    /// Function call locations where this value was inlined from.
    pub inline_locations: Vec<StableLocation<'db>>,
}
impl<'db> Location<'db> {
    pub fn new(stable_location: StableLocation<'db>) -> Self {
        Self { stable_location, notes: vec![], inline_locations: vec![] }
    }

    /// Creates a new Location with the given note as the last note.
    pub fn with_note(mut self, note: DiagnosticNote<'db>) -> Self {
        self.notes.push(note);
        self
    }

    /// Creates a new Location with the given note as the last note.
    pub fn maybe_with_note(mut self, note: Option<DiagnosticNote<'db>>) -> Self {
        let Some(note) = note else {
            return self;
        };
        self.notes.push(note);
        self
    }

    /// Creates a new Location with a note from the given text and location.
    pub fn add_note_with_location(
        self,
        db: &'db dyn Database,
        text: &str,
        location: LocationId<'db>,
    ) -> Self {
        self.with_note(DiagnosticNote::with_location(
            text.into(),
            location.long(db).stable_location.span_in_file(db),
        ))
    }
}

impl<'db> DebugWithDb<'db> for Location<'db> {
    type Db = dyn Database;
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>, db: &'db dyn Database) -> std::fmt::Result {
        self.stable_location.span_in_file(db).fmt(f, db)?;

        for note in &self.notes {
            f.write_str("\nnote: ")?;
            note.fmt(f, db)?;
        }
        Ok(())
    }
}

impl<'db> LocationId<'db> {
    /// Returns the location with the added inlining location of it.
    pub fn inlined(self, db: &'db dyn Database, inlining_location: StableLocation<'db>) -> Self {
        let mut location = self.long(db).clone();
        location.inline_locations.push(inlining_location);
        location.intern(db)
    }
    /// Returns all relevant stable pointers of the location.
    pub fn all_locations(self, db: &'db dyn Database) -> Vec<StableLocation<'db>> {
        let location = self.long(db);
        let mut all_locations = vec![location.stable_location];
        all_locations.extend(location.inline_locations.iter().cloned());
        all_locations
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct VariableMarker;
pub type VariableId = Id<VariableMarker>;
pub type VariableArena<'db> = Arena<Variable<'db>, DefaultArenaBehavior<VariableMarker>>;

/// Represents a usage of a variable.
///
/// For example if we have:
///
/// fn foo(a: u32) {
///     1 + a
/// }
///
/// Then the right hand side of the tail expression `1 + a` is a VarUsage object with
/// the variable id of the variable `a` and the location:
///     1 + a
///         ^
/// Note that the location associated with the variable that was assigned to 'a' is
/// fn foo(a: u32)
///        ^
/// and it is different from the location in the VarUsage.
///
/// The tail expression `1 + a`  is also going to be assigned a variable and a VarUsage.
/// in that case, the location of both the variable and the usage will be the same.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct VarUsage<'db> {
    pub var_id: VariableId,
    pub location: LocationId<'db>,
}

/// A lowered function code using flat blocks.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Lowered<'db> {
    /// Diagnostics produced while lowering.
    pub diagnostics: Diagnostics<'db, LoweringDiagnostic<'db>>,
    /// Function signature.
    pub signature: Signature<'db>,
    /// Arena of allocated lowered variables.
    pub variables: VariableArena<'db>,
    /// Arena of allocated lowered blocks.
    pub blocks: Blocks<'db>,
    /// function parameters, including implicits.
    pub parameters: Vec<VariableId>,
}

unsafe impl<'db> salsa::Update for Lowered<'db> {
    unsafe fn maybe_update(old_pointer: *mut Self, new_value: Self) -> bool {
        let old_value = unsafe { &mut *old_pointer };
        let res = unsafe {
            Diagnostics::maybe_update(&mut old_value.diagnostics, new_value.diagnostics)
                | Signature::maybe_update(&mut old_value.signature, new_value.signature)
        } | (old_value.blocks != new_value.blocks);
        if res {
            old_value.variables = new_value.variables;
            old_value.parameters = new_value.parameters;
            old_value.blocks = new_value.blocks;
        }
        res
    }
}

/// Remapping of lowered variable ids. Useful for convergence of branches.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct VarRemapping<'db> {
    /// Map from new_var to old_var (since new_var cannot appear twice, but old_var can).
    pub remapping: OrderedHashMap<VariableId, VarUsage<'db>>,
}
impl<'db> Deref for VarRemapping<'db> {
    type Target = OrderedHashMap<VariableId, VarUsage<'db>>;

    fn deref(&self) -> &Self::Target {
        &self.remapping
    }
}
impl<'db> DerefMut for VarRemapping<'db> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.remapping
    }
}

/// A block of statements.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Block<'db> {
    /// Statements sequence running one after the other in the block, in a linear flow.
    pub statements: Vec<Statement<'db>>,
    /// Describes how this block ends.
    pub end: BlockEnd<'db>,
}
impl<'db> Default for Block<'db> {
    fn default() -> Self {
        Self { statements: Default::default(), end: BlockEnd::NotSet }
    }
}
impl<'db> Block<'db> {
    pub fn is_set(&self) -> bool {
        !matches!(self.end, BlockEnd::NotSet)
    }
}

/// Describes what happens to the program flow at the end of a [`Block`].
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum BlockEnd<'db> {
    /// The block was created but still needs to be populated. Block must not be in this state in
    /// the end of the lowering phase.
    NotSet,
    /// This block ends with a `return` statement, exiting the function.
    Return(Vec<VarUsage<'db>>, LocationId<'db>),
    /// This block ends with a panic.
    Panic(VarUsage<'db>),
    /// This block ends with a jump to a different block.
    Goto(BlockId, VarRemapping<'db>),
    Match {
        info: MatchInfo<'db>,
    },
}

/// Lowered variable representation.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Variable<'db> {
    /// Semantic type of the variable.
    pub ty: semantic::TypeId<'db>,
    /// Location of the variable.
    pub location: LocationId<'db>,
    /// The semantic type info of the variable.
    pub info: TypeInfo<'db>,
}

impl<'db> DebugWithDb<'db> for Variable<'db> {
    type Db = LoweredFormatter<'db>;

    fn fmt(&self, f: &mut std::fmt::Formatter<'_>, _ctx: &Self::Db) -> std::fmt::Result {
        write!(f, "Variable({:?})", self.ty)
    }
}

impl<'db> Variable<'db> {
    pub fn new(
        db: &'db dyn Database,
        ctx: ImplLookupContextId<'db>,
        ty: semantic::TypeId<'db>,
        location: LocationId<'db>,
    ) -> Self {
        Self { ty, location, info: db.type_info(ctx, ty) }
    }

    /// Returns a new variable with the type, with info calculated with the default context.
    pub fn with_default_context(
        db: &'db dyn Database,
        ty: semantic::TypeId<'db>,
        location: LocationId<'db>,
    ) -> Self {
        Self {
            ty,
            location,
            info: TypeInfo {
                copyable: db.copyable(ty),
                droppable: db.droppable(ty),
                destruct_impl: Err(InferenceError::Ambiguity(Ambiguity::WillNotInfer(
                    concrete_destruct_trait(db, ty),
                ))),
                panic_destruct_impl: Err(InferenceError::Ambiguity(Ambiguity::WillNotInfer(
                    concrete_panic_destruct_trait(db, ty),
                ))),
            },
        }
    }
}

/// Lowered statement.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Statement<'db> {
    // Values.
    Const(StatementConst<'db>),

    // Flow control.
    Call(StatementCall<'db>),

    // Structs (including tuples).
    StructConstruct(StatementStructConstruct<'db>),
    StructDestructure(StatementStructDestructure<'db>),

    // Enums.
    EnumConstruct(StatementEnumConstruct<'db>),

    Snapshot(StatementSnapshot<'db>),
    Desnap(StatementDesnap<'db>),
}
impl<'db> Statement<'db> {
    pub fn inputs(&self) -> &[VarUsage<'db>] {
        match &self {
            Statement::Const(_stmt) => &[],
            Statement::Call(stmt) => stmt.inputs.as_slice(),
            Statement::StructConstruct(stmt) => stmt.inputs.as_slice(),
            Statement::StructDestructure(stmt) => std::slice::from_ref(&stmt.input),
            Statement::EnumConstruct(stmt) => std::slice::from_ref(&stmt.input),
            Statement::Snapshot(stmt) => std::slice::from_ref(&stmt.input),
            Statement::Desnap(stmt) => std::slice::from_ref(&stmt.input),
        }
    }

    pub fn inputs_mut(&mut self) -> &mut [VarUsage<'db>] {
        match self {
            Statement::Const(_stmt) => &mut [],
            Statement::Call(stmt) => stmt.inputs.as_mut_slice(),
            Statement::StructConstruct(stmt) => stmt.inputs.as_mut_slice(),
            Statement::StructDestructure(stmt) => std::slice::from_mut(&mut stmt.input),
            Statement::EnumConstruct(stmt) => std::slice::from_mut(&mut stmt.input),
            Statement::Snapshot(stmt) => std::slice::from_mut(&mut stmt.input),
            Statement::Desnap(stmt) => std::slice::from_mut(&mut stmt.input),
        }
    }

    pub fn outputs(&self) -> &[VariableId] {
        match self {
            Statement::Const(stmt) => std::slice::from_ref(&stmt.output),
            Statement::Call(stmt) => stmt.outputs.as_slice(),
            Statement::StructConstruct(stmt) => std::slice::from_ref(&stmt.output),
            Statement::StructDestructure(stmt) => stmt.outputs.as_slice(),
            Statement::EnumConstruct(stmt) => std::slice::from_ref(&stmt.output),
            Statement::Snapshot(stmt) => stmt.outputs.as_slice(),
            Statement::Desnap(stmt) => std::slice::from_ref(&stmt.output),
        }
    }

    pub fn outputs_mut(&mut self) -> &mut [VariableId] {
        match self {
            Statement::Const(stmt) => std::slice::from_mut(&mut stmt.output),
            Statement::Call(stmt) => stmt.outputs.as_mut_slice(),
            Statement::StructConstruct(stmt) => std::slice::from_mut(&mut stmt.output),
            Statement::StructDestructure(stmt) => stmt.outputs.as_mut_slice(),
            Statement::EnumConstruct(stmt) => std::slice::from_mut(&mut stmt.output),
            Statement::Snapshot(stmt) => stmt.outputs.as_mut_slice(),
            Statement::Desnap(stmt) => std::slice::from_mut(&mut stmt.output),
        }
    }
    pub fn location(&self) -> Option<LocationId<'db>> {
        // TODO(Gil): Add location to all statements.
        match &self {
            Statement::Const(_) => None,
            Statement::Call(stmt) => Some(stmt.location),
            Statement::StructConstruct(_) => None,
            Statement::StructDestructure(stmt) => Some(stmt.input.location),
            Statement::EnumConstruct(stmt) => Some(stmt.input.location),
            Statement::Snapshot(stmt) => Some(stmt.input.location),
            Statement::Desnap(stmt) => Some(stmt.input.location),
        }
    }
    pub fn location_mut(&mut self) -> Option<&mut LocationId<'db>> {
        match self {
            Statement::Const(_) => None,
            Statement::Call(stmt) => Some(&mut stmt.location),
            Statement::StructConstruct(_) => None,
            Statement::StructDestructure(stmt) => Some(&mut stmt.input.location),
            Statement::EnumConstruct(stmt) => Some(&mut stmt.input.location),
            Statement::Snapshot(stmt) => Some(&mut stmt.input.location),
            Statement::Desnap(stmt) => Some(&mut stmt.input.location),
        }
    }
}

/// A statement that binds a const value to a variable.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct StatementConst<'db> {
    /// The value of the const.
    pub value: ConstValueId<'db>,
    /// The variable to bind the value to.
    pub output: VariableId,
    /// Is the const wrapped in a box.
    pub boxed: bool,
}
impl<'db> StatementConst<'db> {
    /// Creates a new const statement, with the option to wrap the value in a box.
    pub fn new(value: ConstValueId<'db>, output: VariableId, boxed: bool) -> Self {
        Self { value, output, boxed }
    }
    /// Creates a new const statement, not boxed.
    pub fn new_flat(value: ConstValueId<'db>, output: VariableId) -> Self {
        Self::new(value, output, false)
    }
    /// Creates a new const statement with the value wrapped in a box.
    pub fn new_boxed(value: ConstValueId<'db>, output: VariableId) -> Self {
        Self::new(value, output, true)
    }
}

/// A statement that calls a user function.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct StatementCall<'db> {
    /// A function to "call".
    pub function: FunctionId<'db>,
    /// Living variables in current scope to move to the function, as arguments.
    pub inputs: Vec<VarUsage<'db>>,
    /// Is the last input a coupon for the function call. See
    /// [semantic::ExprFunctionCall::coupon_arg] for more information.
    pub with_coupon: bool,
    /// New variables to be introduced into the current scope from the function outputs.
    pub outputs: Vec<VariableId>,
    /// Is the call to be inlined as part of the specialization wrapper function.
    pub is_specialization_base_call: bool,
    /// Location for the call.
    pub location: LocationId<'db>,
}

/// A statement that constructs a variant of an enum with a single argument, and binds it to a
/// variable.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct StatementEnumConstruct<'db> {
    pub variant: ConcreteVariant<'db>,
    /// A living variable in current scope to wrap with the variant.
    pub input: VarUsage<'db>,
    /// The variable to bind the value to.
    pub output: VariableId,
}

/// A statement that constructs a struct (tuple included) into a new variable.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct StatementStructConstruct<'db> {
    pub inputs: Vec<VarUsage<'db>>,
    /// The variable to bind the value to.
    pub output: VariableId,
}

/// A statement that destructures a struct (tuple included), introducing its elements as new
/// variables.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct StatementStructDestructure<'db> {
    /// A living variable in current scope to destructure.
    pub input: VarUsage<'db>,
    /// The variables to bind values to.
    pub outputs: Vec<VariableId>,
}

/// A statement that takes a snapshot of a variable.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct StatementSnapshot<'db> {
    pub input: VarUsage<'db>,
    pub outputs: [VariableId; 2],
}
impl<'db> StatementSnapshot<'db> {
    pub fn new(
        input: VarUsage<'db>,
        output_original: VariableId,
        output_snapshot: VariableId,
    ) -> Self {
        Self { input, outputs: [output_original, output_snapshot] }
    }
    pub fn original(&self) -> VariableId {
        self.outputs[0]
    }
    pub fn snapshot(&self) -> VariableId {
        self.outputs[1]
    }
}

/// A statement that desnaps a variable.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct StatementDesnap<'db> {
    pub input: VarUsage<'db>,
    /// The variable to bind the value to.
    pub output: VariableId,
}

/// An arm of a match statement.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MatchArm<'db> {
    /// The selector of the arm.
    pub arm_selector: MatchArmSelector<'db>,

    /// The block_id where the relevant arm is implemented.
    pub block_id: BlockId,

    /// The list of variable ids introduced in this arm.
    pub var_ids: Vec<VariableId>,
}

/// A statement that calls an extern function with branches, and "calls" a possibly different block
/// for each branch.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MatchExternInfo<'db> {
    // TODO(spapini): ConcreteExternFunctionId once it exists.
    /// A concrete external function to call.
    pub function: FunctionId<'db>,
    /// Living variables in current scope to move to the function, as arguments.
    pub inputs: Vec<VarUsage<'db>>,
    /// Match arms. All blocks should have the same rets.
    /// Order must be identical to the order in the definition of the enum.
    pub arms: Vec<MatchArm<'db>>,
    /// Location for the call.
    pub location: LocationId<'db>,
}

/// A statement that matches an enum, and "calls" a possibly different block for each branch.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MatchEnumInfo<'db> {
    pub concrete_enum_id: ConcreteEnumId<'db>,
    /// A living variable in current scope to match on.
    pub input: VarUsage<'db>,
    /// Match arms. All blocks should have the same rets.
    /// Order must be identical to the order in the definition of the enum.
    pub arms: Vec<MatchArm<'db>>,
    /// Location for the match.
    pub location: LocationId<'db>,
}
/// A statement that matches an index enum for matching on felt252, and "calls" a possibly different
/// block for each branch.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MatchEnumValue<'db> {
    pub num_of_arms: usize,

    /// A living variable in current scope to match on.
    pub input: VarUsage<'db>,
    /// Match arms. All blocks should have the same rets.
    pub arms: Vec<MatchArm<'db>>,
    /// Location for the match.
    pub location: LocationId<'db>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum MatchInfo<'db> {
    Enum(MatchEnumInfo<'db>),
    Extern(MatchExternInfo<'db>),
    Value(MatchEnumValue<'db>),
}
impl<'db> MatchInfo<'db> {
    pub fn inputs(&self) -> &[VarUsage<'db>] {
        match self {
            MatchInfo::Enum(s) => std::slice::from_ref(&s.input),
            MatchInfo::Extern(s) => s.inputs.as_slice(),
            MatchInfo::Value(s) => std::slice::from_ref(&s.input),
        }
    }

    pub fn inputs_mut(&mut self) -> &mut [VarUsage<'db>] {
        match self {
            MatchInfo::Enum(s) => std::slice::from_mut(&mut s.input),
            MatchInfo::Extern(s) => s.inputs.as_mut_slice(),
            MatchInfo::Value(s) => std::slice::from_mut(&mut s.input),
        }
    }
    pub fn arms(&self) -> &[MatchArm<'db>] {
        match self {
            MatchInfo::Enum(s) => &s.arms,
            MatchInfo::Extern(s) => &s.arms,
            MatchInfo::Value(s) => &s.arms,
        }
    }
    pub fn arms_mut(&mut self) -> &mut [MatchArm<'db>] {
        match self {
            MatchInfo::Enum(s) => &mut s.arms,
            MatchInfo::Extern(s) => &mut s.arms,
            MatchInfo::Value(s) => &mut s.arms,
        }
    }
    pub fn location(&self) -> &LocationId<'db> {
        match self {
            MatchInfo::Enum(s) => &s.location,
            MatchInfo::Extern(s) => &s.location,
            MatchInfo::Value(s) => &s.location,
        }
    }
    pub fn location_mut(&mut self) -> &mut LocationId<'db> {
        match self {
            MatchInfo::Enum(s) => &mut s.location,
            MatchInfo::Extern(s) => &mut s.location,
            MatchInfo::Value(s) => &mut s.location,
        }
    }
}

/// Used in graph algorithms, and describes how to construct the edges in function dependency graph.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum DependencyType {
    /// A function depends on another function if it may call it.
    Call,
    /// A function depends on another function if its cost depends on the other function's cost.
    Cost,
}

/// The requested lowering stage.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum LoweringStage {
    /// Direct translation from the semantic stage and concretization.
    Monomorphized,
    /// After lowering stages that may change the signature of functions, such as `lower_panics`.
    /// Specifically:
    /// * Adds `withdraw_gas` calls.
    /// * Adds panics.
    /// * Adds destructor calls.
    /// * scrub units.
    PreOptimizations,
    /// Lowering with baseline optimizations - specifically, adds the stages at
    /// `baseline_optimization_strategy`.
    PostBaseline,
    /// Lowering with all of the optimizations - specifically, adds the stages at
    /// `final_optimization_strategy`
    Final,
}