From 84e32147748ce77103022cf7a1ad879d72b83d0f Mon Sep 17 00:00:00 2001 From: Arnab Ghosh Date: Tue, 4 Aug 2026 15:56:59 -0700 Subject: [PATCH] fix: centralize unsafe MLIR access and deny unsafe elsewhere (#415) --- Cargo.toml | 5 +- ...04-quantum-circ-as-unregistered-dialect.md | 7 +- docs/agents/code-quality.md | 4 +- mlir_bridge/src/diagnostics.rs | 26 ++- mlir_bridge/src/ffi.rs | 173 ++++++++++++++++++ mlir_bridge/src/fixed_physical.rs | 22 +-- mlir_bridge/src/lib.rs | 1 + .../src/passes/classical_region_fusion.rs | 21 ++- mlir_bridge/src/passes/clifford_t_opt.rs | 28 ++- .../src/passes/compiler_uncomputation.rs | 21 +-- mlir_bridge/src/passes/depth_scheduling.rs | 34 ++-- mlir_bridge/src/passes/gate_cancellation.rs | 28 ++- .../src/passes/measurement_deferral.rs | 22 ++- mlir_bridge/src/passes/native_gate_decomp.rs | 38 ++-- mlir_bridge/src/passes/rotation_merging.rs | 29 ++- mlir_bridge/src/passes/sabre_routing.rs | 43 ++--- mlir_bridge/src/passes/zx_simplification.rs | 16 +- quonc/src/main.rs | 14 +- 18 files changed, 346 insertions(+), 186 deletions(-) create mode 100644 mlir_bridge/src/ffi.rs diff --git a/Cargo.toml b/Cargo.toml index de571996..9da63f06 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -65,7 +65,10 @@ license-file = "LICENSE" repository = "https://github.com/arniber21/quon" [workspace.lints.rust] -unsafe_code = "allow" +unsafe_code = "deny" + +[workspace.lints.clippy] +undocumented_unsafe_blocks = "warn" [profile.release] lto = true diff --git a/docs/adr/0004-quantum-circ-as-unregistered-dialect.md b/docs/adr/0004-quantum-circ-as-unregistered-dialect.md index 4507740d..204d80df 100644 --- a/docs/adr/0004-quantum-circ-as-unregistered-dialect.md +++ b/docs/adr/0004-quantum-circ-as-unregistered-dialect.md @@ -36,7 +36,10 @@ IRDL, and it would still not give us Rust verifier callbacks. every op builder, so a malformed op cannot be constructed through the public API, and by exposing `verify` for callers that build ops by hand. - Diagnostic *emission* is not wrapped by Melior, so all error reporting flows through the - `diagnostics` module, which isolates the single `unsafe` `mlirEmitError` boundary behind a - `Result`/Writer-style abstraction. + `diagnostics` module, which delegates to the `ffi` module's safe `emit_error` wrapper. The + `ffi` module is the sole audited unsafe boundary in `mlir_bridge`: it owns + `mlirEmitError`, `mlirOperationSetAttributeByName`, `mlirOperationSetOperand`, and + external-pass context lifetime erasure. The workspace denies `unsafe_code` everywhere + else. - Custom types are opaque strings (`!quantum.qubit`, `!quantum.circ`); type checks compare the printed form rather than a registered `TypeID`. diff --git a/docs/agents/code-quality.md b/docs/agents/code-quality.md index 728f031c..632723a7 100644 --- a/docs/agents/code-quality.md +++ b/docs/agents/code-quality.md @@ -76,10 +76,10 @@ When adding fallible parsing or serialization, add at least: happy-path unit tes | **quonc** | `anyhow::Result` | CLI driver; may aggregate errors from the pipeline (`quonc/src/main.rs`). | | **frontend** | `thiserror` (`TypeError`, etc.) | Library API returns typed errors; span-aware reporting via `ariadne`. | | **backend** | `thiserror` (`BackendError`) | Descriptor JSON → domain conversion; every fallible path returns `Result` (`backend/src/error.rs`). | -| **mlir_bridge** | `thiserror` per module + **`Diagnostics` monad** | Verifiers return `Result<(), E>`; passes accumulate with `Diagnostics::report` and flush once at the FFI boundary (`mlir_bridge/src/diagnostics.rs`). | +| **mlir_bridge** | `thiserror` per module + **`Diagnostics` monad** | Verifiers return `Result<(), E>`; passes accumulate with `Diagnostics::report` and flush once at the FFI boundary (`mlir_bridge/src/ffi.rs`). | | **zx** | Typed errors (follow workspace convention) | Graph transforms; no `anyhow` in library code. | -**Diagnostics monad (mlir_bridge):** Dialect verifiers and passes stay pure Rust. They build a `Diagnostics` accumulator, fold `Result` values with `.report(location, result)`, and only `Diagnostics::emit` crosses into unsafe MLIR C API. Do not call `mlirEmitError` outside `diagnostics.rs`. +**Diagnostics monad (mlir_bridge):** Dialect verifiers and passes stay pure Rust. They build a `Diagnostics` accumulator, fold `Result` values with `.report(location, result)`, and only `Diagnostics::emit` flushes to MLIR — via the safe `ffi::emit_error` wrapper. The `ffi` module (`mlir_bridge/src/ffi.rs`) is the sole audited unsafe boundary: it owns `mlirEmitError`, `mlirOperationSetAttributeByName`, `mlirOperationSetOperand`, and external-pass context lifetime erasure (`PassContext`). The workspace denies `unsafe_code` everywhere else. Do not call `mlir-sys` functions outside `ffi.rs`. **anyhow:** Reserved for the **quonc** binary. Library crates should use `thiserror` enums (or plain `Result` with a small `E`). Taskless rule `no-anyhow-in-lib-src` enforces this on new code. diff --git a/mlir_bridge/src/diagnostics.rs b/mlir_bridge/src/diagnostics.rs index b882bef2..37b17669 100644 --- a/mlir_bridge/src/diagnostics.rs +++ b/mlir_bridge/src/diagnostics.rs @@ -9,19 +9,19 @@ //! accumulator (a Writer-style "diagnostic monad") and return it, or fold a //! typed [`Result`] into it with [`Diagnostics::report`]. None of that code //! touches the FFI boundary. -//! * [`Diagnostics::emit`] is the *single* place in the crate that crosses -//! into the unsafe MLIR C API. It is `safe` to call: the one `unsafe` block -//! it contains is self-contained and sound (a valid location plus a -//! NUL-free C string). +//! * [`Diagnostics::emit`] flushes the accumulator to MLIR via the safe +//! [`crate::ffi::emit_error`] wrapper. The `unsafe` `mlirEmitError` call +//! lives solely in [`crate::ffi`], alongside all other unsafe MLIR FFI in +//! the crate. //! -//! This keeps the unsafe surface to one auditable function while the rest of -//! the bridge composes diagnostics with ordinary `Result`/iterator combinators. +//! This keeps the unsafe surface to one audited module ([`crate::ffi`]) while +//! the rest of the bridge composes diagnostics with ordinary +//! `Result`/iterator combinators. use std::ffi::CString; use std::fmt; use melior::ir::Location; -use mlir_sys::mlirEmitError; /// A single error diagnostic anchored at an IR [`Location`]. #[derive(Clone)] @@ -51,9 +51,9 @@ impl<'c> Diagnostic<'c> { /// Emits this diagnostic into MLIR's diagnostic engine. /// - /// This is the only function in the crate that calls into the MLIR C - /// diagnostic API. The message is sanitized of interior NUL bytes so the - /// `CString` conversion is infallible. + /// The message is sanitized of interior NUL bytes so the `CString` + /// conversion is infallible. The actual FFI call is delegated to the safe + /// [`crate::ffi::emit_error`] wrapper — this module contains no `unsafe` code. fn emit(&self) { let sanitized: String = self .message @@ -61,11 +61,7 @@ impl<'c> Diagnostic<'c> { .map(|c| if c == '\0' { '?' } else { c }) .collect(); let message = CString::new(sanitized).unwrap_or_default(); - // SAFETY: `self.location` is a live MLIR location owned by the context, - // and `message` is a valid NUL-terminated C string that outlives the - // call. `mlirEmitError` copies the message and does not retain the - // pointer. - unsafe { mlirEmitError(self.location.to_raw(), message.as_ptr()) }; + crate::ffi::emit_error(&self.location, &message); } } diff --git a/mlir_bridge/src/ffi.rs b/mlir_bridge/src/ffi.rs new file mode 100644 index 00000000..5cf4a814 --- /dev/null +++ b/mlir_bridge/src/ffi.rs @@ -0,0 +1,173 @@ +//! Centralized unsafe FFI boundary for the MLIR bridge. +//! +//! All `unsafe` code in `mlir_bridge` is confined to this module. Pass +//! implementations, verifiers, and the [`crate::diagnostics`] accumulator +//! compose the safe wrappers exported here and never touch raw `mlir-sys` +//! pointers directly. +//! +//! Three concerns are centralized: +//! * Error emission (`mlirEmitError`) — called by [`crate::diagnostics`]. +//! * Raw operation mutation (`mlirOperationSetAttributeByName`, +//! `mlirOperationSetOperand`). +//! * External-pass context lifetime erasure ([`PassContext`] + +//! [`with_context`]). +//! +//! Every `unsafe` block below carries a `SAFETY` comment tied to the upstream +//! FFI contract or the MLIR pass-framework lifetime guarantee. + +#![allow(unsafe_code)] + +use std::ffi::CString; + +use melior::StringRef; +use melior::ir::{Attribute, AttributeLike, Location, OperationRef, Value, ValueLike}; +use melior::{Context, ContextRef}; +use mlir_sys::{ + MlirContext, mlirEmitError, mlirOperationSetAttributeByName, mlirOperationSetOperand, +}; + +// ─── Error emission ───────────────────────────────────────────────────────── + +/// Emits an error diagnostic at `location` with the given C string message. +/// +/// This is the sole call site for `mlirEmitError` in the crate. The message +/// must already be a valid NUL-terminated C string (the [`crate::diagnostics`] +/// module sanitizes interior NULs before calling). +pub(crate) fn emit_error(location: &Location<'_>, message: &CString) { + // SAFETY: `location` is a live MLIR location backed by the context that + // owns it. `message` is a valid NUL-terminated C string whose backing + // buffer outlives the call. `mlirEmitError` copies the message internally + // and does not retain the pointer after returning. + unsafe { mlirEmitError(location.to_raw(), message.as_ptr()) }; +} + +// ─── Raw operation mutation ───────────────────────────────────────────────── + +/// Sets a named attribute on an MLIR operation. +/// +/// Safe wrapper for `mlirOperationSetAttributeByName`. The operation and +/// attribute must belong to the same context. +pub(crate) fn set_operation_attribute<'c>( + op: OperationRef<'c, '_>, + name: &str, + attribute: &Attribute<'c>, +) { + // SAFETY: `op` is a live operation reference. `name` is borrowed as a + // `StringRef` for the duration of the call; the C function copies it + // internally via `MlirStringRef` and does not retain the pointer. + // `attribute` is a live attribute owned by the same context. The function + // performs an in-place attribute update and does not retain any pointer + // after returning. + unsafe { + mlirOperationSetAttributeByName( + op.to_raw(), + StringRef::new(name).to_raw(), + attribute.to_raw(), + ); + } +} + +/// Replaces a single operand of an MLIR operation. +/// +/// Safe wrapper for `mlirOperationSetOperand`. The operation and value must +/// belong to the same context. The caller is responsible for ensuring `index` +/// is a valid operand slot. +pub(crate) fn set_operation_operand<'c, 'a>( + op: OperationRef<'c, 'a>, + index: isize, + value: &Value<'c, 'a>, +) { + // SAFETY: `op` is a live operation reference. `index` is a valid operand + // index validated by the caller. `value` is a live SSA value belonging to + // the same context. The function updates the operand slot in place and + // does not retain the pointers after returning. + unsafe { + mlirOperationSetOperand(op.to_raw(), index, value.to_raw()); + } +} + +// ─── External-pass context lifetime erasure ──────────────────────────────── + +/// Erases the external-pass context lifetime so that a `'static` pass struct +/// can hold a context reference between `initialize` and `run`. +/// +/// MLIR's `RunExternalPass` trait hands passes a `ContextRef<'c>` in +/// `initialize`, but the pass struct must be `'static` and cannot hold a +/// `&'c Context` directly. The pass framework guarantees that the context +/// remains valid for the entire lifetime of the pass. +/// +/// `PassContext` stores the raw `MlirContext` handle (a `Copy` value, not a +/// pointer to stack memory) extracted in `initialize`. In `run`, the handle is +/// reconstructed into a `ContextRef` whose lifetime is scoped to the +/// [`with_context`] closure, ensuring the `&Context` is always valid. +/// +/// # Why not store a `&Context` pointer? +/// +/// `ContextRef::to_ref` uses `transmute` to return a `&Context` that points to +/// the `ContextRef` itself (they share the same layout: a single +/// `MlirContext`). Storing that pointer for later use is undefined behaviour +/// because the `ContextRef` is a stack local. Storing the `MlirContext` handle +/// by value and reconstructing the `ContextRef` in `run` avoids this pitfall. +#[derive(Clone, Copy, Default)] +pub(crate) struct PassContext { + raw: Option, +} + +impl PassContext { + /// Creates an empty context store (no context captured yet). + pub(crate) fn new() -> Self { + Self { raw: None } + } + + /// Captures the MLIR context handle for later retrieval. + /// + /// The `MlirContext` handle is extracted by value from the `ContextRef`. + /// No reference to the `ContextRef` (which is a stack local) is retained. + pub(crate) fn capture<'c>(&mut self, context: ContextRef<'c>) { + // SAFETY: `to_ref` transmutes `&ContextRef` into `&Context` (same + // layout), and `to_raw` copies the `MlirContext` handle out by value. + // The `&Context` is alive only for this expression; the stored + // `MlirContext` is an independent value, not a dangling pointer. + self.raw = Some(unsafe { context.to_ref().to_raw() }); + } + + /// Whether a context has been captured. + pub(crate) fn is_captured(&self) -> bool { + self.raw.is_some() + } + + /// Returns the raw `MlirContext` handle, or `None` if [`capture`](Self::capture) + /// was never called. + /// + /// The handle is safe to pass to [`with_context`]. + pub(crate) fn raw(&self) -> Option { + self.raw + } +} + +/// Runs a closure with a borrowed `&'c Context` reconstructed from a stored +/// `MlirContext` handle. +/// +/// The `ContextRef` is created as a local variable inside this function, so +/// the `&Context` returned by `to_ref` (which `transmute`s `&ContextRef` into +/// `&Context`) is valid for the duration of the closure call. After the +/// closure returns, both are dropped — no dangling pointer escapes. +/// +/// # Panics +/// +/// Panics if `raw` is a null/invalid handle. The pass framework guarantees +/// validity, so this should never happen in practice. +pub(crate) fn with_context<'c, R>(raw: MlirContext, f: impl FnOnce(&'c Context) -> R) -> R { + // SAFETY: The `MlirContext` handle was obtained in `PassContext::capture` + // from a `ContextRef` provided by `RunExternalPass::initialize`. The pass + // framework guarantees the context outlives every `run` call, so the + // handle is still valid here. + let context_ref = unsafe { ContextRef::from_raw(raw) }; + // SAFETY: `context_ref` is a local variable alive for the duration of this + // function. `to_ref` transmutes `&context_ref` into `&'c Context` — the + // reference points to `context_ref`'s stack slot, which is valid until + // this function returns. The closure consumes the reference before that, + // so no dangling pointer escapes. + let context = unsafe { context_ref.to_ref() }; + f(context) +} diff --git a/mlir_bridge/src/fixed_physical.rs b/mlir_bridge/src/fixed_physical.rs index 77f7b5de..99e4328c 100644 --- a/mlir_bridge/src/fixed_physical.rs +++ b/mlir_bridge/src/fixed_physical.rs @@ -29,15 +29,14 @@ use backend::BackendTarget; use melior::Context; -use melior::StringRef; use melior::ir::Module; use melior::ir::attribute::IntegerAttribute; use melior::ir::operation::OperationLike; use melior::ir::r#type::IntegerType; -use melior::ir::{AttributeLike, BlockLike, OperationRef, RegionLike}; -use mlir_sys::mlirOperationSetAttributeByName; +use melior::ir::{BlockLike, OperationRef, RegionLike}; use crate::dialect::quantum_dynamic; +use crate::ffi; use crate::metrics; use crate::passes::{ depth_scheduling, native_gate_decomp, sabre_routing, sabre_routing::SabreCost, @@ -85,8 +84,7 @@ pub fn run_fixed_physical( pub fn corrupt_phys_qubit_attrs(context: &Context, module: &Module<'_>, bogus: i32) { let attr: melior::ir::Attribute<'_> = IntegerAttribute::new(IntegerType::new(context, 32).into(), i64::from(bogus)).into(); - let raw = attr.to_raw(); - let name = StringRef::new(quantum_dynamic::attr::PHYS_QUBIT).to_raw(); + let name = quantum_dynamic::attr::PHYS_QUBIT; let Some(body) = module .as_operation() .region(0) @@ -95,13 +93,13 @@ pub fn corrupt_phys_qubit_attrs(context: &Context, module: &Module<'_>, bogus: i else { return; }; - corrupt_block(body, name, raw); + corrupt_block(body, name, &attr); } fn corrupt_block<'c, 'a>( block: melior::ir::BlockRef<'c, 'a>, - name: mlir_sys::MlirStringRef, - attr: mlir_sys::MlirAttribute, + name: &str, + attr: &melior::ir::Attribute<'c>, ) { let mut op = block.first_operation(); while let Some(current) = op { @@ -112,15 +110,13 @@ fn corrupt_block<'c, 'a>( fn corrupt_op<'c, 'a>( op: OperationRef<'c, 'a>, - name: mlir_sys::MlirStringRef, - attr: mlir_sys::MlirAttribute, + name: &str, + attr: &melior::ir::Attribute<'c>, ) { // Only ops that already have a phys_qubit attr are touched — overwriting // a non-existent attr would *add* one, which is not the test's intent. if op.attribute(quantum_dynamic::attr::PHYS_QUBIT).is_ok() { - unsafe { - mlirOperationSetAttributeByName(op.to_raw(), name, attr); - } + ffi::set_operation_attribute(op, name, attr); } let count = op.region_count(); for index in 0..count { diff --git a/mlir_bridge/src/lib.rs b/mlir_bridge/src/lib.rs index 5b5306e1..eddf6b8e 100644 --- a/mlir_bridge/src/lib.rs +++ b/mlir_bridge/src/lib.rs @@ -2,6 +2,7 @@ pub mod circ_extract; pub mod diagnostics; +pub mod ffi; pub mod dialect; pub mod dynamic_walk; pub mod emit; diff --git a/mlir_bridge/src/passes/classical_region_fusion.rs b/mlir_bridge/src/passes/classical_region_fusion.rs index b0ecf68f..63bff212 100644 --- a/mlir_bridge/src/passes/classical_region_fusion.rs +++ b/mlir_bridge/src/passes/classical_region_fusion.rs @@ -30,6 +30,7 @@ use melior::pass::{ExternalPass, Pass, RunExternalPass, create_external}; use melior::{Context, ContextRef, IrRewriter}; use thiserror::Error; +use crate::ffi::PassContext; use crate::diagnostics::Diagnostics; use crate::dialect::{quantum_circ, quantum_dynamic}; use crate::passes::qubit_wiring::{self, WireTracker}; @@ -688,29 +689,31 @@ static CLASSICAL_REGION_FUSION_PASS_ID: PassId = PassId; #[derive(Clone)] struct ClassicalRegionFusion { - context: usize, + context: PassContext, } impl ClassicalRegionFusion { fn new() -> Self { - Self { context: 0 } + Self { context: PassContext::new() } } } impl<'c> RunExternalPass<'c> for ClassicalRegionFusion { fn initialize(&mut self, context: ContextRef<'c>) { - self.context = unsafe { context.to_ref() as *const Context as usize }; + self.context.capture(context); } fn run(&mut self, operation: OperationRef<'c, '_>, pass: ExternalPass<'_>) { - if self.context == 0 { + let Some(raw) = self.context.raw() else { pass.signal_failure(); return; - } - let context = unsafe { &*(self.context as *const Context) }; - let mut diagnostics = Diagnostics::new(); - fuse_module(context, operation, &mut diagnostics); - if !diagnostics.emit() { + }; + let success = crate::ffi::with_context(raw, |context| { + let mut diagnostics = Diagnostics::new(); + fuse_module(context, operation, &mut diagnostics); + diagnostics.emit() + }); + if !success { pass.signal_failure(); } } diff --git a/mlir_bridge/src/passes/clifford_t_opt.rs b/mlir_bridge/src/passes/clifford_t_opt.rs index 05e8382a..c2aa204c 100644 --- a/mlir_bridge/src/passes/clifford_t_opt.rs +++ b/mlir_bridge/src/passes/clifford_t_opt.rs @@ -31,16 +31,15 @@ use std::collections::HashMap; -use melior::StringRef; use melior::ir::attribute::{BoolAttribute, StringAttribute}; use melior::ir::operation::OperationLike; use melior::ir::r#type::TypeId; -use melior::ir::{Attribute, AttributeLike, BlockLike, OperationRef, RegionLike, Value, ValueLike}; +use melior::ir::{Attribute, BlockLike, OperationRef, RegionLike, Value, ValueLike}; use melior::pass::{ExternalPass, Pass, RunExternalPass, create_external}; use melior::{Context, ContextRef, IrRewriter}; -use mlir_sys::mlirOperationSetAttributeByName; use quon_core::DepthExpr; +use crate::ffi::{self, PassContext}; use crate::dialect::quantum_circ::{self, attr}; use crate::passes::{phase_polynomial, stabilizer_tableau}; @@ -84,13 +83,7 @@ fn read_depth_attr<'c: 'a, 'a, O: OperationLike<'c, 'a>>(operation: &O) -> Depth fn set_func_depth<'c, 'a>(context: &'c Context, func: OperationRef<'c, 'a>, depth: &DepthExpr) { let attribute: Attribute<'c> = StringAttribute::new(context, &depth.to_sexpr()).into(); - unsafe { - mlirOperationSetAttributeByName( - func.to_raw(), - StringRef::new(attr::DEPTH).to_raw(), - attribute.to_raw(), - ); - } + ffi::set_operation_attribute(func, attr::DEPTH, &attribute); } fn gate_is_clifford(name: &str) -> bool { @@ -311,27 +304,28 @@ static CLIFFORD_T_OPT_PASS_ID: PassId = PassId; #[derive(Clone)] struct CliffordTOpt { - context: usize, + context: PassContext, } impl CliffordTOpt { fn new() -> Self { - Self { context: 0 } + Self { context: PassContext::new() } } } impl<'c> RunExternalPass<'c> for CliffordTOpt { fn initialize(&mut self, context: ContextRef<'c>) { - self.context = unsafe { context.to_ref() as *const Context as usize }; + self.context.capture(context); } fn run(&mut self, operation: OperationRef<'c, '_>, pass: ExternalPass<'_>) { - if self.context == 0 { + let Some(raw) = self.context.raw() else { pass.signal_failure(); return; - } - let context = unsafe { &*(self.context as *const Context) }; - optimize_module(context, operation); + }; + crate::ffi::with_context(raw, |context| { + optimize_module(context, operation); + }); } } diff --git a/mlir_bridge/src/passes/compiler_uncomputation.rs b/mlir_bridge/src/passes/compiler_uncomputation.rs index fbc0f6fb..893c9d75 100644 --- a/mlir_bridge/src/passes/compiler_uncomputation.rs +++ b/mlir_bridge/src/passes/compiler_uncomputation.rs @@ -10,9 +10,9 @@ use melior::ir::r#type::TypeId; use melior::ir::{BlockLike, OperationRef, RegionLike, Value, ValueLike}; use melior::pass::{ExternalPass, Pass, RunExternalPass, create_external}; use melior::{Context, ContextRef, IrRewriter}; -use mlir_sys::mlirOperationSetOperand; use crate::dialect::quantum_circ::{self, attr}; +use crate::ffi::{self, PassContext}; #[derive(Clone)] struct RecordedGate { @@ -56,9 +56,7 @@ fn inverse_name(name: &str) -> Option { fn set_return_operands<'c, 'a>(return_op: OperationRef<'c, 'a>, wires: &[Value<'c, 'a>]) { for (index, value) in wires.iter().enumerate() { - unsafe { - mlirOperationSetOperand(return_op.to_raw(), index as isize, value.to_raw()); - } + ffi::set_operation_operand(return_op, index as isize, value); } } @@ -192,26 +190,27 @@ static COMPILER_UNCOMPUTATION_PASS_ID: PassId = PassId; #[derive(Clone)] struct CompilerUncomputation { - context: usize, + context: PassContext, } impl CompilerUncomputation { fn new() -> Self { - Self { context: 0 } + Self { context: PassContext::new() } } } impl<'c> RunExternalPass<'c> for CompilerUncomputation { fn initialize(&mut self, context: ContextRef<'c>) { - self.context = unsafe { context.to_ref() as *const Context as usize }; + self.context.capture(context); } fn run(&mut self, operation: OperationRef<'c, '_>, _pass: ExternalPass<'_>) { - if self.context == 0 { + let Some(raw) = self.context.raw() else { return; - } - let context = unsafe { &*(self.context as *const Context) }; - uncompute_module(context, operation); + }; + crate::ffi::with_context(raw, |context| { + uncompute_module(context, operation); + }); } } diff --git a/mlir_bridge/src/passes/depth_scheduling.rs b/mlir_bridge/src/passes/depth_scheduling.rs index 102ca577..0b8a4506 100644 --- a/mlir_bridge/src/passes/depth_scheduling.rs +++ b/mlir_bridge/src/passes/depth_scheduling.rs @@ -12,11 +12,11 @@ use backend::target::BackendTarget; use melior::ir::attribute::IntegerAttribute; use melior::ir::operation::OperationLike; use melior::ir::r#type::TypeId; -use melior::ir::{AttributeLike, BlockLike, OperationRef, RegionLike}; +use melior::ir::{BlockLike, OperationRef, RegionLike}; use melior::pass::{ExternalPass, Pass, RunExternalPass, create_external}; -use melior::{Context, ContextRef, StringRef}; -use mlir_sys::mlirOperationSetAttributeByName; +use melior::{Context, ContextRef}; +use crate::ffi::{self, PassContext}; use crate::dialect::quantum_circ; use crate::dynamic_walk::{self, DynamicVisitor}; @@ -43,13 +43,7 @@ fn set_schedule_time<'c>(context: &'c Context, op: OperationRef<'c, '_>, time: i time, ) .into(); - unsafe { - mlirOperationSetAttributeByName( - op.to_raw(), - StringRef::new("schedule_time").to_raw(), - attribute.to_raw(), - ); - } + ffi::set_operation_attribute(op, "schedule_time", &attribute); } struct GateStep<'c, 'a> { @@ -259,14 +253,14 @@ static DEPTH_SCHEDULING_PASS_ID: PassId = PassId; #[derive(Clone)] struct DepthScheduling { - context: usize, + context: PassContext, target: Arc, } impl DepthScheduling { fn new(target: BackendTarget) -> Self { Self { - context: 0, + context: PassContext::new(), target: Arc::new(target), } } @@ -274,19 +268,21 @@ impl DepthScheduling { impl<'c> RunExternalPass<'c> for DepthScheduling { fn initialize(&mut self, context: ContextRef<'c>) { - self.context = unsafe { context.to_ref() as *const Context as usize }; + self.context.capture(context); } fn run(&mut self, operation: OperationRef<'c, '_>, pass: ExternalPass<'_>) { - if self.context == 0 { + let Some(raw) = self.context.raw() else { pass.signal_failure(); return; - } - let Some(target) = self.target.fixed_target() else { - return; }; - let context = unsafe { &*(self.context as *const Context) }; - schedule_module(context, target, operation); + let target = self.target.clone(); + crate::ffi::with_context(raw, |context| { + let Some(target) = target.fixed_target() else { + return; + }; + schedule_module(context, target, operation); + }); } } diff --git a/mlir_bridge/src/passes/gate_cancellation.rs b/mlir_bridge/src/passes/gate_cancellation.rs index 6d3eac9d..d6c616d3 100644 --- a/mlir_bridge/src/passes/gate_cancellation.rs +++ b/mlir_bridge/src/passes/gate_cancellation.rs @@ -5,16 +5,15 @@ use std::collections::HashMap; -use melior::StringRef; use melior::ir::attribute::{IntegerAttribute, StringAttribute}; use melior::ir::operation::OperationLike; use melior::ir::r#type::TypeId; -use melior::ir::{Attribute, AttributeLike, BlockLike, OperationRef, RegionLike, Value, ValueLike}; +use melior::ir::{Attribute, BlockLike, OperationRef, RegionLike, Value, ValueLike}; use melior::pass::{ExternalPass, Pass, RunExternalPass, create_external}; use melior::{Context, ContextRef, IrRewriter}; -use mlir_sys::mlirOperationSetAttributeByName; use quon_core::DepthExpr; +use crate::ffi::{self, PassContext}; use crate::dialect::{ quantum_circ::{self, attr}, quantum_dynamic, @@ -135,13 +134,7 @@ fn cancel_pair<'c, 'a>( fn set_func_depth<'c, 'a>(context: &'c Context, func: OperationRef<'c, 'a>, depth: &DepthExpr) { let attribute: Attribute<'c> = StringAttribute::new(context, &depth.to_sexpr()).into(); - unsafe { - mlirOperationSetAttributeByName( - func.to_raw(), - StringRef::new(attr::DEPTH).to_raw(), - attribute.to_raw(), - ); - } + ffi::set_operation_attribute(func, attr::DEPTH, &attribute); } fn cancel_in_block<'c, 'a>(context: &'c Context, block: melior::ir::BlockRef<'c, 'a>) -> i64 { @@ -273,27 +266,28 @@ static GATE_CANCELLATION_PASS_ID: PassId = PassId; #[derive(Clone)] struct GateCancellation { - context: usize, + context: PassContext, } impl GateCancellation { fn new() -> Self { - Self { context: 0 } + Self { context: PassContext::new() } } } impl<'c> RunExternalPass<'c> for GateCancellation { fn initialize(&mut self, context: ContextRef<'c>) { - self.context = unsafe { context.to_ref() as *const Context as usize }; + self.context.capture(context); } fn run(&mut self, operation: OperationRef<'c, '_>, pass: ExternalPass<'_>) { - if self.context == 0 { + let Some(raw) = self.context.raw() else { pass.signal_failure(); return; - } - let context = unsafe { &*(self.context as *const Context) }; - cancel_module(context, operation); + }; + crate::ffi::with_context(raw, |context| { + cancel_module(context, operation); + }); } } diff --git a/mlir_bridge/src/passes/measurement_deferral.rs b/mlir_bridge/src/passes/measurement_deferral.rs index 288166da..5796ac55 100644 --- a/mlir_bridge/src/passes/measurement_deferral.rs +++ b/mlir_bridge/src/passes/measurement_deferral.rs @@ -14,6 +14,8 @@ use melior::{Context, ContextRef, IrRewriter}; use quon_core::DepthExpr; use thiserror::Error; +use crate::ffi::PassContext; + use crate::diagnostics::Diagnostics; use crate::dialect::{quantum_circ, quantum_dynamic}; use crate::passes::qubit_wiring::{self, WireTracker}; @@ -424,29 +426,31 @@ static MEASUREMENT_DEFERRAL_PASS_ID: PassId = PassId; #[derive(Clone)] struct MeasurementDeferral { - context: usize, + context: PassContext, } impl MeasurementDeferral { fn new() -> Self { - Self { context: 0 } + Self { context: PassContext::new() } } } impl<'c> RunExternalPass<'c> for MeasurementDeferral { fn initialize(&mut self, context: ContextRef<'c>) { - self.context = unsafe { context.to_ref() as *const Context as usize }; + self.context.capture(context); } fn run(&mut self, operation: OperationRef<'c, '_>, pass: ExternalPass<'_>) { - if self.context == 0 { + let Some(raw) = self.context.raw() else { pass.signal_failure(); return; - } - let context = unsafe { &*(self.context as *const Context) }; - let mut diagnostics = Diagnostics::new(); - defer_module(context, operation, &mut diagnostics); - if !diagnostics.emit() { + }; + let success = crate::ffi::with_context(raw, |context| { + let mut diagnostics = Diagnostics::new(); + defer_module(context, operation, &mut diagnostics); + diagnostics.emit() + }); + if !success { pass.signal_failure(); } } diff --git a/mlir_bridge/src/passes/native_gate_decomp.rs b/mlir_bridge/src/passes/native_gate_decomp.rs index 149c3408..2c5e6770 100644 --- a/mlir_bridge/src/passes/native_gate_decomp.rs +++ b/mlir_bridge/src/passes/native_gate_decomp.rs @@ -8,17 +8,16 @@ use std::sync::Arc; use backend::decompose::{decompose_named_single, decompose_named_two}; use backend::target::{BackendTarget, FixedTarget}; -use melior::StringRef; use melior::ir::attribute::{BoolAttribute, FloatAttribute, IntegerAttribute, StringAttribute}; use melior::ir::operation::OperationLike; use melior::ir::r#type::TypeId; -use melior::ir::{AttributeLike, BlockLike, Location, OperationRef, RegionLike, Value, ValueLike}; +use melior::ir::{BlockLike, Location, OperationRef, RegionLike, Value, ValueLike}; use melior::pass::{ExternalPass, Pass, RunExternalPass, create_external}; use melior::{Context, ContextRef, IrRewriter}; -use mlir_sys::mlirOperationSetAttributeByName; use thiserror::Error; use crate::diagnostics::Diagnostics; +use crate::ffi::{self, PassContext}; use crate::dialect::quantum_circ::{self, attr}; use crate::dialect::quantum_dynamic; @@ -70,13 +69,7 @@ fn native_gate_names(target: &FixedTarget) -> Vec { fn set_native_gate<'c>(context: &'c Context, op: OperationRef<'c, '_>, native: bool) { let attribute: melior::ir::Attribute<'_> = BoolAttribute::new(context, native).into(); - unsafe { - mlirOperationSetAttributeByName( - op.to_raw(), - StringRef::new("native_gate").to_raw(), - attribute.to_raw(), - ); - } + ffi::set_operation_attribute(op, "native_gate", &attribute); } #[allow(clippy::too_many_arguments)] @@ -335,14 +328,14 @@ static NATIVE_GATE_DECOMP_PASS_ID: PassId = PassId; #[derive(Clone)] struct NativeGateDecomp { - context: usize, + context: PassContext, target: Arc, } impl NativeGateDecomp { fn new(target: BackendTarget) -> Self { Self { - context: 0, + context: PassContext::new(), target: Arc::new(target), } } @@ -350,20 +343,23 @@ impl NativeGateDecomp { impl<'c> RunExternalPass<'c> for NativeGateDecomp { fn initialize(&mut self, context: ContextRef<'c>) { - self.context = unsafe { context.to_ref() as *const Context as usize }; + self.context.capture(context); } fn run(&mut self, operation: OperationRef<'c, '_>, pass: ExternalPass<'_>) { - if self.context == 0 { + let Some(raw) = self.context.raw() else { pass.signal_failure(); return; - } - let context = unsafe { &*(self.context as *const Context) }; - let mut diagnostics = Diagnostics::new(); - if let Some(target) = self.target.fixed_target() { - decompose_module(context, target, operation, &mut diagnostics); - } - if !diagnostics.emit() { + }; + let target = self.target.clone(); + let success = crate::ffi::with_context(raw, |context| { + let mut diagnostics = Diagnostics::new(); + if let Some(t) = target.fixed_target() { + decompose_module(context, t, operation, &mut diagnostics); + } + diagnostics.emit() + }); + if !success { pass.signal_failure(); } } diff --git a/mlir_bridge/src/passes/rotation_merging.rs b/mlir_bridge/src/passes/rotation_merging.rs index f3f9fdbf..572b1f5b 100644 --- a/mlir_bridge/src/passes/rotation_merging.rs +++ b/mlir_bridge/src/passes/rotation_merging.rs @@ -6,12 +6,12 @@ use std::f64::consts::TAU; use melior::ir::attribute::{FloatAttribute, IntegerAttribute, StringAttribute}; use melior::ir::operation::OperationLike; use melior::ir::r#type::TypeId; -use melior::ir::{Attribute, AttributeLike, BlockLike, OperationRef, RegionLike, Value, ValueLike}; +use melior::ir::{Attribute, BlockLike, OperationRef, RegionLike, Value, ValueLike}; use melior::pass::{ExternalPass, Pass, RunExternalPass, create_external}; -use melior::{Context, ContextRef, IrRewriter, StringRef}; -use mlir_sys::mlirOperationSetAttributeByName; +use melior::{Context, ContextRef, IrRewriter}; use quon_core::DepthExpr; +use crate::ffi::{self, PassContext}; use crate::dialect::{ quantum_circ::{self, attr}, quantum_dynamic, @@ -131,13 +131,7 @@ fn parse_rotation<'c, 'a>( fn set_func_depth<'c, 'a>(context: &'c Context, func: OperationRef<'c, 'a>, depth: &DepthExpr) { let attribute: Attribute<'c> = StringAttribute::new(context, &depth.to_sexpr()).into(); - unsafe { - mlirOperationSetAttributeByName( - func.to_raw(), - StringRef::new(attr::DEPTH).to_raw(), - attribute.to_raw(), - ); - } + ffi::set_operation_attribute(func, attr::DEPTH, &attribute); } fn merge_pair<'c, 'a>( @@ -334,27 +328,28 @@ static ROTATION_MERGING_PASS_ID: PassId = PassId; #[derive(Clone)] struct RotationMerging { - context: usize, + context: PassContext, } impl RotationMerging { fn new() -> Self { - Self { context: 0 } + Self { context: PassContext::new() } } } impl<'c> RunExternalPass<'c> for RotationMerging { fn initialize(&mut self, context: ContextRef<'c>) { - self.context = unsafe { context.to_ref() as *const Context as usize }; + self.context.capture(context); } fn run(&mut self, operation: OperationRef<'c, '_>, pass: ExternalPass<'_>) { - if self.context == 0 { + let Some(raw) = self.context.raw() else { pass.signal_failure(); return; - } - let context = unsafe { &*(self.context as *const Context) }; - merge_module(context, operation); + }; + crate::ffi::with_context(raw, |context| { + merge_module(context, operation); + }); } } diff --git a/mlir_bridge/src/passes/sabre_routing.rs b/mlir_bridge/src/passes/sabre_routing.rs index b634616f..bfcc8107 100644 --- a/mlir_bridge/src/passes/sabre_routing.rs +++ b/mlir_bridge/src/passes/sabre_routing.rs @@ -7,17 +7,16 @@ use std::collections::HashMap; use std::sync::Arc; use backend::target::{BackendTarget, FixedTarget}; -use melior::StringRef; use melior::ir::attribute::IntegerAttribute; use melior::ir::operation::OperationLike; use melior::ir::r#type::TypeId; -use melior::ir::{AttributeLike, BlockLike, Location, OperationRef, RegionLike, Value, ValueLike}; +use melior::ir::{BlockLike, Location, OperationRef, RegionLike, Value, ValueLike}; use melior::pass::{ExternalPass, Pass, RunExternalPass, create_external}; use melior::{Context, ContextRef}; -use mlir_sys::{mlirOperationSetAttributeByName, mlirOperationSetOperand}; use thiserror::Error; use crate::diagnostics::Diagnostics; +use crate::ffi::{self, PassContext}; use crate::dialect::{quantum_circ, quantum_dynamic}; use crate::passes::qubit_wiring::{self, WireTracker}; @@ -27,13 +26,7 @@ fn set_i32_attr<'c>(context: &'c Context, op: OperationRef<'c, '_>, key: &str, v i64::from(value), ) .into(); - unsafe { - mlirOperationSetAttributeByName( - op.to_raw(), - StringRef::new(key).to_raw(), - attribute.to_raw(), - ); - } + ffi::set_operation_attribute(op, key, &attribute); } #[derive(Clone, Copy, Debug)] pub struct SabreCost { @@ -188,9 +181,7 @@ fn set_qubit_operands<'c, 'a>(gate: OperationRef<'c, 'a>, values: &[Value<'c, 'a continue; } if let Some(value) = values.get(qubit_index) { - unsafe { - mlirOperationSetOperand(gate.to_raw(), operand_index as isize, value.to_raw()); - } + ffi::set_operation_operand(gate, operand_index as isize, value); } qubit_index += 1; } @@ -740,7 +731,7 @@ static SABRE_ROUTING_PASS_ID: PassId = PassId; #[derive(Clone)] struct SabreRouting { - context: usize, + context: PassContext, target: Arc, cost: SabreCost, } @@ -748,7 +739,7 @@ struct SabreRouting { impl SabreRouting { fn new(target: BackendTarget, cost: SabreCost) -> Self { Self { - context: 0, + context: PassContext::new(), target: Arc::new(target), cost, } @@ -757,20 +748,24 @@ impl SabreRouting { impl<'c> RunExternalPass<'c> for SabreRouting { fn initialize(&mut self, context: ContextRef<'c>) { - self.context = unsafe { context.to_ref() as *const Context as usize }; + self.context.capture(context); } fn run(&mut self, operation: OperationRef<'c, '_>, pass: ExternalPass<'_>) { - if self.context == 0 { + let Some(raw) = self.context.raw() else { pass.signal_failure(); return; - } - let context = unsafe { &*(self.context as *const Context) }; - let mut diagnostics = Diagnostics::new(); - if let Some(target) = self.target.fixed_target() { - route_module(context, target, self.cost, operation, &mut diagnostics); - } - if !diagnostics.emit() { + }; + let target = self.target.clone(); + let cost = self.cost; + let success = crate::ffi::with_context(raw, |context| { + let mut diagnostics = Diagnostics::new(); + if let Some(t) = target.fixed_target() { + route_module(context, t, cost, operation, &mut diagnostics); + } + diagnostics.emit() + }); + if !success { pass.signal_failure(); } } diff --git a/mlir_bridge/src/passes/zx_simplification.rs b/mlir_bridge/src/passes/zx_simplification.rs index 418504df..76381a83 100644 --- a/mlir_bridge/src/passes/zx_simplification.rs +++ b/mlir_bridge/src/passes/zx_simplification.rs @@ -11,6 +11,7 @@ use melior::pass::{ExternalPass, Pass, RunExternalPass, create_external}; use melior::{Context, ContextRef}; use zx::{GateRef, circuit_to_zx, simplify, zx_to_circuit}; +use crate::ffi::PassContext; use crate::circ_extract; use crate::dialect::{quantum_circ, quantum_dynamic}; @@ -141,26 +142,27 @@ static ZX_SIMPLIFICATION_PASS_ID: PassId = PassId; #[derive(Clone)] struct ZxSimplification { - context: usize, + context: PassContext, } impl ZxSimplification { fn new() -> Self { - Self { context: 0 } + Self { context: PassContext::new() } } } impl<'c> RunExternalPass<'c> for ZxSimplification { fn initialize(&mut self, context: ContextRef<'c>) { - self.context = unsafe { context.to_ref() as *const Context as usize }; + self.context.capture(context); } fn run(&mut self, operation: OperationRef<'c, '_>, _pass: ExternalPass<'_>) { - if self.context == 0 { + let Some(raw) = self.context.raw() else { return; - } - let context = unsafe { &*(self.context as *const Context) }; - simplify_module(context, operation); + }; + crate::ffi::with_context(raw, |context| { + simplify_module(context, operation); + }); } } diff --git a/quonc/src/main.rs b/quonc/src/main.rs index 3f34a5cc..8cfe26a8 100644 --- a/quonc/src/main.rs +++ b/quonc/src/main.rs @@ -626,11 +626,21 @@ fn run() -> Result { Ok(ExitCode::SUCCESS) } +/// Propagates `--color` to `QUONC_COLOR` so our own stderr styling honors it. +/// +/// `std::env::set_var` is `unsafe` in Edition 2024 because it is not re-entrant +/// with concurrent `std::env::var` calls. This CLI sets the variable early in +/// `main`, before any worker threads exist, so the data race cannot occur. +#[allow(unsafe_code)] fn apply_color_env(cli: &Cli) { - // clap ColorChoice is set at parse time via attribute; also honor --color for - // our own stderr styling via QUONC_COLOR. + // SAFETY: This runs single-threaded during CLI argument processing, before + // any tokio runtime or background thread is spawned. No concurrent + // `std::env::var` reader exists at this point, so the mutation is free of + // data races. match cli.color { + // SAFETY: see function-level comment — single-threaded, pre-runtime. CliColor::Always => unsafe { std::env::set_var("QUONC_COLOR", "always") }, + // SAFETY: see function-level comment — single-threaded, pre-runtime. CliColor::Never => unsafe { std::env::set_var("QUONC_COLOR", "never") }, CliColor::Auto => {} }