diff --git a/Cargo.lock b/Cargo.lock index 3b7af70..84ad91c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -21,14 +21,15 @@ dependencies = [ [[package]] name = "ahash" -version = "0.8.12" +version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +checksum = "e89da841a80418a9b391ebaea17f5c112ffaaa96f621d2c285b5174da76b9011" dependencies = [ "cfg-if", + "getrandom 0.2.17", "once_cell", "version_check", - "zerocopy", + "zerocopy 0.7.35", ] [[package]] @@ -40,6 +41,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + [[package]] name = "anstream" version = "1.0.0" @@ -204,6 +211,12 @@ version = "3.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" +[[package]] +name = "bytemuck" +version = "1.25.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" + [[package]] name = "bytes" version = "1.11.1" @@ -299,6 +312,20 @@ version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f88a43d011fc4a6876cb7344703e297c71dda42494fee094d5f7c76bf13f746" +[[package]] +name = "cobra-mba" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78abaa2dea834236a4affc21235fa43fe92622a4bc735923cf551c604bce848f" +dependencies = [ + "ahash", + "hashbrown 0.14.5", + "indexmap", + "smallvec", + "thiserror 1.0.69", + "wide", +] + [[package]] name = "colorchoice" version = "1.0.5" @@ -692,6 +719,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" dependencies = [ "ahash", + "allocator-api2", ] [[package]] @@ -1470,7 +1498,7 @@ version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" dependencies = [ - "zerocopy", + "zerocopy 0.8.48", ] [[package]] @@ -1793,6 +1821,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "safe_arch" +version = "0.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96b02de82ddbe1b636e6170c21be622223aea188ef2e139be0a5b219ec215323" +dependencies = [ + "bytemuck", +] + [[package]] name = "same-file" version = "1.0.6" @@ -1945,6 +1982,7 @@ version = "0.1.0" dependencies = [ "binbit", "blake3", + "cobra-mba", "qfbvsmtrs", "rumba-core", "rusqlite", @@ -2574,6 +2612,16 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "wide" +version = "0.7.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ce5da8ecb62bcd8ec8b7ea19f69a51275e91299be594ea5cc6ef7819e16cd03" +dependencies = [ + "bytemuck", + "safe_arch", +] + [[package]] name = "winapi-util" version = "0.1.11" @@ -2906,13 +2954,33 @@ dependencies = [ "zip", ] +[[package]] +name = "zerocopy" +version = "0.7.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" +dependencies = [ + "zerocopy-derive 0.7.35", +] + [[package]] name = "zerocopy" version = "0.8.48" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eed437bf9d6692032087e337407a86f04cd8d6a16a37199ed57949d415bd68e9" dependencies = [ - "zerocopy-derive", + "zerocopy-derive 0.8.48", +] + +[[package]] +name = "zerocopy-derive" +version = "0.7.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 2fbdaf0..6c859aa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,6 +24,8 @@ smt-wire = { path = "crates/smt-wire" } # gh-release feature downloads a matching prebuilt solver for each. binbit = { git = "https://github.com/bint-disasm/binbit", rev = "b4996145c06f39ce5bbc79a959975d7d20b6cf5e" } blake3 = "1" +# The crates.io package is `cobra-mba`; the library itself is named `cobra`. +cobra = { package = "cobra-mba", version = "0.3.0" } dashu = { version = "0.4.2", features = ["num-traits"] } rumba-core = "1.0.0" rusqlite = { version = "0.32", features = ["bundled"] } diff --git a/README.md b/README.md index 705ce7a..9d3c00d 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ It targets binary analysis, lifting, symbolic execution, and IR experiments wher ## Backends - Solve and optimize: [`z3`](https://docs.rs/z3/latest/z3/), [`binbit`](https://github.com/bint-disasm/binbit), and the standalone `qfbvsmtrs` crate. -- Simplify: [Rumba](https://github.com/thalium/rumba) for supported 64-bit-or-smaller MBA expression islands. +- Simplify: [Rumba](https://github.com/thalium/rumba) and [CoBRA](https://github.com/binsnake/cobra) run as a chain over supported 64-bit-or-smaller MBA expression islands; CoBRA adopts only Lean-certified rewrites by default. - Text compatibility: SMT-LIB `QF_BV` scripts are parsed into the same binary IR used by binary clients. The binary protocol is the main API. SMT-LIB support exists for tooling compatibility and test reuse. @@ -225,7 +225,7 @@ See `docs/qfbvsmtrs-validation.md` for maintainer validation gates and corpus-ru - `crates/smt-wire` — Rust high-level client API plus server-side wire-format internals and validators. - `crates/qfbvsmtrs` — standalone pure-Rust `QF_BV` bit-blasting solver crate and CLI. -- `crates/smt-server` — TCP server, Rumba simplifier integration, solver backend integration, SMT-LIB frontend. +- `crates/smt-server` — TCP server, Rumba and CoBRA simplifier integrations, solver backend integration, SMT-LIB frontend. - `python` — Python client package and tests. - `cpp` — C++17 header-only package, CMake target, and tests. - `docs/architecture.md` — current crate/server/backend architecture. diff --git a/crates/smt-qfbv-smtlib/src/lib.rs b/crates/smt-qfbv-smtlib/src/lib.rs index 0de22a9..52e2603 100644 --- a/crates/smt-qfbv-smtlib/src/lib.rs +++ b/crates/smt-qfbv-smtlib/src/lib.rs @@ -527,7 +527,7 @@ pub fn parse_annotation(items: &[SExpr]) -> Result().0 { let key = atom(&pair[0])?; if !key.starts_with(':') { return Err(FrontendError::invalid( diff --git a/crates/smt-server/Cargo.toml b/crates/smt-server/Cargo.toml index ad9b377..6718776 100644 --- a/crates/smt-server/Cargo.toml +++ b/crates/smt-server/Cargo.toml @@ -7,6 +7,7 @@ license.workspace = true [dependencies] binbit.workspace = true +cobra.workspace = true qfbvsmtrs = { workspace = true, features = ["wire"] } rumba-core.workspace = true smt-qfbv-smtlib.workspace = true diff --git a/crates/smt-server/src/cobra_backend.rs b/crates/smt-server/src/cobra_backend.rs new file mode 100644 index 0000000..6bbbe46 --- /dev/null +++ b/crates/smt-server/src/cobra_backend.rs @@ -0,0 +1,695 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use cobra::{ + simplify_expr, Expr as CobraExpr, Kind as CobraKind, Options as CobraOptions, ProofLevel, + SimplifyOutcomeKind, +}; +use smt_wire::raw::{ + tag, BinaryRequest, Command, ExprBuilder, ExprView, NodeRef, RawNode, SimplifyBlock, WireError, +}; + +use crate::backend::{Backend, QueryResult, SolveContext}; + +const MAX_COBRA_WIDTH: u32 = 64; + +/// Simplifier backend built on [CoBRA](https://github.com/binsnake/cobra) +/// (`cobra-mba`), a worklist-driven mixed Boolean-arithmetic simplifier. +/// +/// Like the Rumba backend, CoBRA only understands pure same-width bit-vector +/// MBA trees, so this backend rewrites each maximal MBA island of the target +/// expression in place and copies the structural skeleton verbatim. +/// +/// CoBRA verifies candidates with finite probing plus a Lean certificate +/// replay. In the default certified mode only rewrites backed by a replayable +/// Lean certificate are adopted; islands without one are left unchanged. The +/// spot-checked mode also adopts probe-verified rewrites, which raises the +/// simplification rate dramatically but can in principle return an expression +/// that differs from the input at an unprobed point, so it should only be used +/// on non-adversarial inputs. +#[derive(Debug, Clone)] +pub struct CobraBackend { + require_certificate: bool, +} + +impl CobraBackend { + /// Certified mode: only adopt rewrites CoBRA backs with a Lean certificate. + pub fn new() -> Self { + Self { + require_certificate: true, + } + } + + /// Spot-checked mode: also adopt probe-verified rewrites without a + /// certificate. Higher simplification rate, probabilistic soundness. + pub fn spot_checked() -> Self { + Self { + require_certificate: false, + } + } +} + +impl Default for CobraBackend { + fn default() -> Self { + Self::new() + } +} + +impl Backend for CobraBackend { + fn name(&self) -> &'static str { + "cobra" + } + + fn handle(&self, request: &BinaryRequest) -> smt_wire::Result { + match request.envelope.command { + Command::Simplify => { + simplify_request(request, self.require_certificate).map(QueryResult::simplified) + } + Command::Solve | Command::Minimize | Command::Maximize => { + Ok(QueryResult::unknown("cobra only supports SIMPLIFY")) + } + } + } + + fn handle_with_context( + &self, + request: &BinaryRequest, + context: &SolveContext, + ) -> smt_wire::Result { + if context.is_cancelled() { + return Ok(QueryResult::unknown("cobra request cancelled before start")); + } + self.handle(request) + } +} + +/// Simplify the target expression by rewriting each maximal MBA island in +/// place, exactly like the Rumba backend: walk the DAG, copy the structural +/// skeleton verbatim, and hand every maximal MBA island to CoBRA -- abstracting +/// any non-MBA / wrong-width child of an island as a fresh opaque variable that +/// maps back to the recursively-simplified subtree. +fn simplify_request( + request: &BinaryRequest, + require_certificate: bool, +) -> smt_wire::Result { + let target = request + .target_ref() + .ok_or_else(|| WireError::invalid("simplify request", "missing target_node"))?; + let view = request.expression_view()?; + let mut simplifier = IslandSimplifier::new(view, require_certificate); + match simplifier.process(target) { + Ok(root) => Ok(SimplifyBlock { + expression: simplifier.builder.to_bytes()?, + target_node: root, + }), + // Any structural failure falls back to returning the input unchanged, which is + // always a sound (if unhelpful) simplification result. + Err(_) => Ok(identity_simplify(request, target)), + } +} + +fn identity_simplify(request: &BinaryRequest, target_node: NodeRef) -> SimplifyBlock { + SimplifyBlock { + expression: request.expression.clone(), + target_node, + } +} + +/// True for the bit-vector operators CoBRA understands; these are the roots of +/// MBA islands. Leaves (VAR/CONST) are copied directly; everything else is +/// structural. `BV_LSHR` by a constant amount also converts inside an island, +/// but a lone shift is not worth an island of its own. +fn is_mba_op(tag: u8) -> bool { + matches!( + tag, + tag::BV_NOT + | tag::BV_NEG + | tag::BV_AND + | tag::BV_OR + | tag::BV_XOR + | tag::BV_ADD + | tag::BV_SUB + | tag::BV_MUL + ) +} + +/// Total island-conversion node visits allowed per request, as a multiple of +/// the input's node count (plus a floor so small requests get full attempts). +/// +/// A declined maximal island is retried on each of its subtrees, and CoBRA's +/// tree conversion re-expands shared wire DAG nodes on every visit, so without +/// a shared cap an adversarial unsimplifiable expression could turn one +/// bounded request into quadratic (or, through sharing, exponential) +/// conversion work and unbounded repeat runs of CoBRA's search pipeline. Once +/// the budget is spent, remaining islands are copied structurally. +const ISLAND_BUDGET_FACTOR: usize = 8; +const ISLAND_BUDGET_FLOOR: usize = 1 << 14; + +struct IslandSimplifier<'a> { + view: ExprView<'a>, + builder: ExprBuilder, + memo: HashMap, + require_certificate: bool, + conversion_budget: usize, +} + +impl<'a> IslandSimplifier<'a> { + fn new(view: ExprView<'a>, require_certificate: bool) -> Self { + let conversion_budget = (view.node_count() as usize) + .saturating_mul(ISLAND_BUDGET_FACTOR) + .max(ISLAND_BUDGET_FLOOR); + Self { + view, + builder: ExprBuilder::new(), + memo: HashMap::new(), + require_certificate, + conversion_budget, + } + } + + fn process(&mut self, reference: NodeRef) -> smt_wire::Result { + if let Some(&cached) = self.memo.get(&reference) { + return Ok(cached); + } + let result = self.process_uncached(reference)?; + self.memo.insert(reference, result); + Ok(result) + } + + fn process_uncached(&mut self, reference: NodeRef) -> smt_wire::Result { + let node = self.view.node(reference.index())?; + if reference.is_bv() + && is_mba_op(node.tag) + && (1..=MAX_COBRA_WIDTH).contains(&node.width) + && self.conversion_budget > 0 + { + if let Some(simplified) = self.try_island(reference, &node)? { + return Ok(simplified); + } + } + self.copy_node(reference, &node) + } + + /// Convert the MBA island rooted at `reference`, simplify it with CoBRA, + /// and lower the result back into the builder. Returns `None` (so the + /// caller copies the node structurally instead) if CoBRA declines the + /// island or fails to improve it under the configured proof requirement. + fn try_island( + &mut self, + reference: NodeRef, + node: &RawNode, + ) -> smt_wire::Result> { + let width = node.width; + let mut conversion = IslandConversion::new(self.view, width, &mut self.conversion_budget); + // A conversion failure (typically the shared budget running out) means + // this island is skipped, not that the request is bad: the caller + // copies the node structurally and any real structural defect + // resurfaces there. + let Ok(expr) = conversion.convert(reference) else { + return Ok(None); + }; + let vars = conversion.vars; + let names = island_var_names(&vars); + let options = CobraOptions { + bitwidth: width, + require_lean_certificate: self.require_certificate, + ..CobraOptions::default() + }; + // A `CobraError` (too many variables, node budget, ...) means CoBRA + // could not handle this island, not that the request is bad: leave the + // island alone and let the caller copy it structurally. + let outcome = match simplify_expr(&expr, &names, options) { + Ok(outcome) => outcome, + Err(_) => return Ok(None), + }; + if outcome.kind != SimplifyOutcomeKind::Simplified { + return Ok(None); + } + // CoBRA's own certificate gate covers the main pipeline, but a few + // side paths still return probe-verified results; in certified mode + // adopt nothing weaker than a Lean-certified rewrite. + if self.require_certificate && outcome.proof_level != ProofLevel::LeanCertified { + return Ok(None); + } + let Some(simplified) = outcome.expr else { + return Ok(None); + }; + // The islands handed over are uniform-width, so a result using + // width-changing operators (or unknown variables) is out of contract. + if !is_lowerable(&simplified, vars.len()) { + return Ok(None); + } + let root = self.lower(&simplified, width, &vars)?; + Ok(Some(root)) + } + + /// Lower a simplified CoBRA expression into the builder; opaque boundary + /// variables resolve to the recursively-simplified subtrees they + /// abstracted. + fn lower( + &mut self, + expr: &CobraExpr, + width: u32, + vars: &[IslandVar], + ) -> smt_wire::Result { + let mask = mask_for_width(width); + match expr.kind { + CobraKind::Variable(index) => match vars.get(index as usize) { + Some(IslandVar::Named(name)) => self.builder.bv_var(name, width), + Some(IslandVar::Boundary(reference)) => self.process(*reference), + None => Err(WireError::invalid( + "cobra expression", + format!("unknown CoBRA variable v{index}"), + )), + }, + CobraKind::Constant(value) => self.builder.bv_const(value & mask, width), + CobraKind::Not => { + let child = self.lower_child(expr, 0, width, vars)?; + self.builder.bv_not(child) + } + CobraKind::Neg => { + let child = self.lower_child(expr, 0, width, vars)?; + self.builder.bv_neg(child) + } + CobraKind::Shr(amount) => { + // CoBRA pins the shift amount in the node; an amount at or + // beyond the width yields zero, matching SMT `bvlshr`. + if u64::from(amount) >= u64::from(width) { + self.builder.bv_const(0, width) + } else { + let child = self.lower_child(expr, 0, width, vars)?; + let amount = self.builder.bv_const(u64::from(amount), width)?; + self.builder.bv_lshr(child, amount) + } + } + CobraKind::Add => self.lower_binary(expr, width, vars, ExprBuilder::bv_add), + CobraKind::Mul => self.lower_binary(expr, width, vars, ExprBuilder::bv_mul), + CobraKind::And => self.lower_binary(expr, width, vars, ExprBuilder::bv_and), + CobraKind::Or => self.lower_binary(expr, width, vars, ExprBuilder::bv_or), + CobraKind::Xor => self.lower_binary(expr, width, vars, ExprBuilder::bv_xor), + CobraKind::ZExt(_) | CobraKind::SExt(_) | CobraKind::Trunc(_) | CobraKind::Concat => { + Err(WireError::invalid( + "cobra expression", + "width-changing operator in uniform-width island result", + )) + } + } + } + + fn lower_child( + &mut self, + expr: &CobraExpr, + index: usize, + width: u32, + vars: &[IslandVar], + ) -> smt_wire::Result { + let child = expr.children.get(index).ok_or_else(|| { + WireError::invalid("cobra expression", "operator node missing a child") + })?; + self.lower(child, width, vars) + } + + fn lower_binary( + &mut self, + expr: &CobraExpr, + width: u32, + vars: &[IslandVar], + op: fn(&mut ExprBuilder, NodeRef, NodeRef) -> smt_wire::Result, + ) -> smt_wire::Result { + let a = self.lower_child(expr, 0, width, vars)?; + let b = self.lower_child(expr, 1, width, vars)?; + op(&mut self.builder, a, b) + } + + fn src_child(&self, node: &RawNode, offset: u32) -> smt_wire::Result { + let index = node + .children + .checked_add(offset) + .ok_or(WireError::IntegerOverflow("child array index"))?; + self.view.child_ref(index) + } + + fn processed_child(&mut self, node: &RawNode, offset: u32) -> smt_wire::Result { + let child = self.src_child(node, offset)?; + self.process(child) + } + + /// Rebuild a node verbatim with its children recursively simplified. + fn copy_node(&mut self, _reference: NodeRef, node: &RawNode) -> smt_wire::Result { + let width = node.width; + match node.tag { + tag::BV_VAR => { + let name = self + .view + .blob_str(node.blob_ref(), "BV variable")? + .to_owned(); + self.builder.bv_var(&name, width) + } + tag::BV_CONST => { + if width <= MAX_COBRA_WIDTH { + self.builder + .bv_const(node.payload & mask_for_width(width), width) + } else { + let bytes = self.view.blob_ref(node.blob_ref())?.to_vec(); + self.builder.bv_const_wide(&bytes, width) + } + } + tag::BV_NOT => { + let child = self.processed_child(node, 0)?; + self.builder.bv_not(child) + } + tag::BV_NEG => { + let child = self.processed_child(node, 0)?; + self.builder.bv_neg(child) + } + tag::BV_AND => self.copy_nary(node, ExprBuilder::bv_and), + tag::BV_OR => self.copy_nary(node, ExprBuilder::bv_or), + tag::BV_XOR => self.copy_nary(node, ExprBuilder::bv_xor), + tag::BV_ADD => self.copy_nary(node, ExprBuilder::bv_add), + tag::BV_MUL => self.copy_nary(node, ExprBuilder::bv_mul), + tag::BV_SUB => self.copy_binary(node, ExprBuilder::bv_sub), + tag::BV_UDIV => self.copy_binary(node, ExprBuilder::bv_udiv), + tag::BV_UREM => self.copy_binary(node, ExprBuilder::bv_urem), + tag::BV_SDIV => self.copy_binary(node, ExprBuilder::bv_sdiv), + tag::BV_SREM => self.copy_binary(node, ExprBuilder::bv_srem), + tag::BV_SMOD => self.copy_binary(node, ExprBuilder::bv_smod), + tag::BV_SHL => self.copy_binary(node, ExprBuilder::bv_shl), + tag::BV_LSHR => self.copy_binary(node, ExprBuilder::bv_lshr), + tag::BV_ASHR => self.copy_binary(node, ExprBuilder::bv_ashr), + tag::BV_EXTRACT => { + let child = self.processed_child(node, 0)?; + self.builder + .bv_extract(child, node.aux_hi as u32, node.aux_lo) + } + tag::BV_CONCAT => { + let high = self.processed_child(node, 0)?; + let low = self.processed_child(node, 1)?; + self.builder.bv_concat(high, low) + } + tag::BV_ZEXT => { + let child = self.processed_child(node, 0)?; + self.builder.bv_zext(child, node.aux_hi) + } + tag::BV_SEXT => { + let child = self.processed_child(node, 0)?; + self.builder.bv_sext(child, node.aux_hi) + } + tag::BV_ITE => { + let cond = self.processed_child(node, 0)?; + let then_value = self.processed_child(node, 1)?; + let else_value = self.processed_child(node, 2)?; + self.builder.bv_ite(cond, then_value, else_value) + } + tag::BV_SELECT => self.copy_select(node), + tag::BOOL_TRUE => self.builder.bool_true(), + tag::BOOL_FALSE => self.builder.bool_false(), + tag::BOOL_VAR => { + let name = self + .view + .blob_str(node.blob_ref(), "Bool variable")? + .to_owned(); + self.builder.bool_var(&name) + } + tag::BOOL_NOT => { + let child = self.processed_child(node, 0)?; + self.builder.bool_not(child) + } + tag::BOOL_AND => self.copy_binary(node, ExprBuilder::bool_and), + tag::BOOL_OR => self.copy_binary(node, ExprBuilder::bool_or), + tag::BOOL_IMPLIES => self.copy_binary(node, ExprBuilder::bool_implies), + tag::BV_EQ => self.copy_binary(node, ExprBuilder::bv_eq), + tag::BV_ULT => self.copy_binary(node, ExprBuilder::bv_ult), + tag::BV_ULE => self.copy_binary(node, ExprBuilder::bv_ule), + tag::BV_SLT => self.copy_binary(node, ExprBuilder::bv_slt), + tag::BV_SLE => self.copy_binary(node, ExprBuilder::bv_sle), + tag::UADD_OVF => self.copy_binary(node, ExprBuilder::uadd_ovf), + tag::SADD_OVF => self.copy_binary(node, ExprBuilder::sadd_ovf), + tag::USUB_OVF => self.copy_binary(node, ExprBuilder::usub_ovf), + tag::SSUB_OVF => self.copy_binary(node, ExprBuilder::ssub_ovf), + tag::UMUL_OVF => self.copy_binary(node, ExprBuilder::umul_ovf), + tag::SMUL_OVF => self.copy_binary(node, ExprBuilder::smul_ovf), + tag::NEG_OVF => { + let child = self.processed_child(node, 0)?; + self.builder.neg_ovf(child) + } + tag::SDIV_OVF => self.copy_binary(node, ExprBuilder::sdiv_ovf), + other => Err(WireError::invalid( + "simplify", + format!("unsupported node tag {other}"), + )), + } + } + + fn copy_nary( + &mut self, + node: &RawNode, + op: fn(&mut ExprBuilder, NodeRef, NodeRef) -> smt_wire::Result, + ) -> smt_wire::Result { + let arity = node.arity as u32; + if arity == 0 { + return Err(WireError::invalid( + "simplify", + "n-ary node with no children", + )); + } + let mut acc = self.processed_child(node, 0)?; + for offset in 1..arity { + let child = self.processed_child(node, offset)?; + acc = op(&mut self.builder, acc, child)?; + } + Ok(acc) + } + + fn copy_binary( + &mut self, + node: &RawNode, + op: fn(&mut ExprBuilder, NodeRef, NodeRef) -> smt_wire::Result, + ) -> smt_wire::Result { + let a = self.processed_child(node, 0)?; + let b = self.processed_child(node, 1)?; + op(&mut self.builder, a, b) + } + + fn copy_select(&mut self, node: &RawNode) -> smt_wire::Result { + let pairs = node.aux_hi as u32; + let mut selectors = Vec::with_capacity(pairs as usize); + let mut values = Vec::with_capacity(pairs as usize); + for pair in 0..pairs { + selectors.push(self.processed_child(node, 2 * pair)?); + values.push(self.processed_child(node, 2 * pair + 1)?); + } + let default = self.processed_child(node, 2 * pairs)?; + self.builder.bv_select(&selectors, &values, default) + } +} + +/// True when `expr` only uses operators and variable indices this backend can +/// lower back into a uniform-width wire island. +fn is_lowerable(expr: &CobraExpr, var_count: usize) -> bool { + match expr.kind { + CobraKind::Variable(index) => (index as usize) < var_count, + CobraKind::Constant(_) => true, + CobraKind::Add + | CobraKind::Mul + | CobraKind::And + | CobraKind::Or + | CobraKind::Xor + | CobraKind::Not + | CobraKind::Neg + | CobraKind::Shr(_) => expr + .children + .iter() + .all(|child| is_lowerable(child, var_count)), + CobraKind::ZExt(_) | CobraKind::SExt(_) | CobraKind::Trunc(_) | CobraKind::Concat => false, + } +} + +/// A leaf of a CoBRA MBA island: either a real bit-vector variable, or an +/// opaque boundary standing in for a non-MBA / wrong-width subtree that will be +/// simplified recursively and spliced back when the island is lowered. +enum IslandVar { + Named(String), + Boundary(NodeRef), +} + +/// CoBRA identifies variables by index but requires unique names alongside. +/// Real variables keep their wire names; boundaries get synthetic names, +/// disambiguated against everything else with underscore suffixes. +fn island_var_names(vars: &[IslandVar]) -> Vec { + let mut used: std::collections::HashSet = vars + .iter() + .filter_map(|var| match var { + IslandVar::Named(name) => Some(name.clone()), + IslandVar::Boundary(_) => None, + }) + .collect(); + vars.iter() + .enumerate() + .map(|(index, var)| match var { + IslandVar::Named(name) => name.clone(), + IslandVar::Boundary(_) => { + let mut name = format!("__cobra_boundary{index}"); + while !used.insert(name.clone()) { + name.push('_'); + } + name + } + }) + .collect() +} + +struct IslandConversion<'a, 'b> { + view: ExprView<'a>, + width: u32, + vars: Vec, + by_name: HashMap, + by_ref: HashMap, + /// Node-visit budget shared across every island attempt of one request; + /// see [`ISLAND_BUDGET_FACTOR`]. + budget: &'b mut usize, +} + +impl<'a, 'b> IslandConversion<'a, 'b> { + fn new(view: ExprView<'a>, width: u32, budget: &'b mut usize) -> Self { + Self { + view, + width, + vars: Vec::new(), + by_name: HashMap::new(), + by_ref: HashMap::new(), + budget, + } + } + + /// Convert a subtree to a CoBRA expression at the island's width. MBA + /// operators of the island width recurse; anything else (other ops, a + /// different width, a wide constant, a Bool) becomes an opaque boundary + /// variable. + /// + /// Conversion never abandons an island for being too large: CoBRA enforces + /// its own variable and node budgets in `simplify_expr` and the resulting + /// error makes the caller copy the island structurally. + fn convert(&mut self, reference: NodeRef) -> smt_wire::Result> { + *self.budget = self.budget.checked_sub(1).ok_or_else(|| { + WireError::invalid("simplify", "cobra island conversion budget exhausted") + })?; + if !reference.is_bv() { + return Ok(self.boundary(reference)); + } + let node = self.view.node(reference.index())?; + if node.width != self.width { + return Ok(self.boundary(reference)); + } + Ok(match node.tag { + tag::BV_VAR => self.named(reference, &node)?, + tag::BV_CONST if node.width <= MAX_COBRA_WIDTH => { + CobraExpr::constant(node.payload & mask_for_width(self.width)) + } + tag::BV_NOT => CobraExpr::not(self.convert_child(&node, 0)?), + tag::BV_NEG => CobraExpr::neg(self.convert_child(&node, 0)?), + tag::BV_AND => self.fold_nary(&node, CobraExpr::and)?, + tag::BV_OR => self.fold_nary(&node, CobraExpr::or)?, + tag::BV_XOR => self.fold_nary(&node, CobraExpr::xor)?, + tag::BV_ADD => self.fold_nary(&node, CobraExpr::add)?, + tag::BV_MUL => self.fold_nary(&node, CobraExpr::mul)?, + tag::BV_SUB => { + let a = self.convert_child(&node, 0)?; + let b = self.convert_child(&node, 1)?; + CobraExpr::add(a, CobraExpr::neg(b)) + } + // A logical shift by a constant amount stays inside the island; + // CoBRA pins the amount into the node. Shifts by an expression + // become boundaries. + tag::BV_LSHR => match self.constant_shift_amount(&node)? { + Some(amount) => CobraExpr::shr(self.convert_child(&node, 0)?, amount), + None => self.boundary(reference), + }, + _ => self.boundary(reference), + }) + } + + /// The shift amount when child 1 is a narrow constant, `None` otherwise. + fn constant_shift_amount(&self, node: &RawNode) -> smt_wire::Result> { + let amount_ref = self.src_child(node, 1)?; + if !amount_ref.is_bv() { + return Ok(None); + } + let amount = self.view.node(amount_ref.index())?; + if amount.tag != tag::BV_CONST || amount.width > MAX_COBRA_WIDTH { + return Ok(None); + } + Ok(Some(amount.payload & mask_for_width(amount.width))) + } + + fn convert_child(&mut self, node: &RawNode, offset: u32) -> smt_wire::Result> { + let child = self.src_child(node, offset)?; + self.convert(child) + } + + /// Fold a wire n-ary node into CoBRA's binary tree shape. + fn fold_nary( + &mut self, + node: &RawNode, + op: fn(Arc, Arc) -> Arc, + ) -> smt_wire::Result> { + let arity = node.arity as u32; + if arity == 0 { + return Err(WireError::invalid( + "simplify", + "n-ary node with no children", + )); + } + let mut acc = self.convert_child(node, 0)?; + for offset in 1..arity { + let child = self.convert_child(node, offset)?; + acc = op(acc, child); + } + Ok(acc) + } + + fn named(&mut self, reference: NodeRef, node: &RawNode) -> smt_wire::Result> { + if let Some(&id) = self.by_ref.get(&reference) { + return Ok(CobraExpr::variable(id as u32)); + } + let name = self + .view + .blob_str(node.blob_ref(), "BV variable")? + .to_owned(); + if let Some(&id) = self.by_name.get(&name) { + self.by_ref.insert(reference, id); + return Ok(CobraExpr::variable(id as u32)); + } + let id = self.vars.len(); + self.vars.push(IslandVar::Named(name.clone())); + self.by_ref.insert(reference, id); + self.by_name.insert(name, id); + Ok(CobraExpr::variable(id as u32)) + } + + fn boundary(&mut self, reference: NodeRef) -> Arc { + if let Some(&id) = self.by_ref.get(&reference) { + return CobraExpr::variable(id as u32); + } + let id = self.vars.len(); + self.vars.push(IslandVar::Boundary(reference)); + self.by_ref.insert(reference, id); + CobraExpr::variable(id as u32) + } + + fn src_child(&self, node: &RawNode, offset: u32) -> smt_wire::Result { + let index = node + .children + .checked_add(offset) + .ok_or(WireError::IntegerOverflow("child array index"))?; + self.view.child_ref(index) + } +} + +fn mask_for_width(width: u32) -> u64 { + debug_assert!((1..=MAX_COBRA_WIDTH).contains(&width)); + if width >= 64 { + u64::MAX + } else { + (1u64 << width) - 1 + } +} diff --git a/crates/smt-server/src/lib.rs b/crates/smt-server/src/lib.rs index 87410ce..efa1ad7 100644 --- a/crates/smt-server/src/lib.rs +++ b/crates/smt-server/src/lib.rs @@ -7,6 +7,7 @@ pub mod backend; pub mod binbit_backend; pub mod cache; +pub mod cobra_backend; pub mod command_router; pub mod pool; pub mod protocol; @@ -15,6 +16,7 @@ pub mod racing; pub(crate) mod recording; pub mod rumba_backend; pub mod server; +pub mod simplify_chain; pub mod smt2; pub mod smtlib; pub mod z3_backend; @@ -22,6 +24,7 @@ pub mod z3_backend; pub use backend::{Backend, CancellationToken, QueryResult, QueryStatus, SolveContext}; pub use binbit_backend::BinbitBackend; pub use cache::{cache_key_for_payload, rebind_cached_response, CacheStats, ResponseCache}; +pub use cobra_backend::CobraBackend; pub use command_router::CommandRouterBackend; pub use pool::PooledBackend; pub use protocol::{handle_binary_frame, handle_binary_request, response_from_query_result}; @@ -34,6 +37,7 @@ pub use recording::{ }; pub use rumba_backend::RumbaBackend; pub use server::{dispatch_payload, dispatch_payload_with_cache, serve_tcp, ServerConfig}; +pub use simplify_chain::SimplifyChainBackend; pub use smt2::{request_to_smt2, Smt2Script, Smt2Variable}; pub use smtlib::{handle_text_frame, parse_smtlib_script, TextQuery, WireSmtLibSink}; pub use z3_backend::Z3Backend; diff --git a/crates/smt-server/src/main.rs b/crates/smt-server/src/main.rs index 3df090b..fee5529 100644 --- a/crates/smt-server/src/main.rs +++ b/crates/smt-server/src/main.rs @@ -3,8 +3,8 @@ use std::sync::Arc; use smt_server::{ default_legacy_recording_tree, migrate_recording_tree, recording_db_path, serve_tcp, - BinbitBackend, CommandRouterBackend, QfbvsmtrsBackend, RacingBackend, RumbaBackend, - ServerConfig, Z3Backend, + BinbitBackend, CobraBackend, CommandRouterBackend, QfbvsmtrsBackend, RacingBackend, + RumbaBackend, ServerConfig, SimplifyChainBackend, Z3Backend, }; fn main() -> std::io::Result<()> { @@ -23,9 +23,13 @@ fn main() -> std::io::Result<()> { ]) .with_default_budget_ms(30_000), ); - let backend = Arc::new(CommandRouterBackend::new(Arc::new(RumbaBackend), solver)); + let simplifier = Arc::new(SimplifyChainBackend::new(vec![ + Arc::new(RumbaBackend), + Arc::new(CobraBackend::default()), + ])); + let backend = Arc::new(CommandRouterBackend::new(simplifier, solver)); eprintln!( - "smt-server listening on {addr} with rumba simplifier + racing solver (z3 crate + binbit + qfbvsmtrs)" + "smt-server listening on {addr} with rumba + cobra simplifier chain + racing solver (z3 crate + binbit + qfbvsmtrs)" ); serve_tcp(addr, ServerConfig::new(backend)) } diff --git a/crates/smt-server/src/simplify_chain.rs b/crates/smt-server/src/simplify_chain.rs new file mode 100644 index 0000000..2e8bba9 --- /dev/null +++ b/crates/smt-server/src/simplify_chain.rs @@ -0,0 +1,125 @@ +use std::sync::Arc; + +use smt_wire::raw::{BinaryRequest, Command, SimplifyBlock, WireError}; + +use crate::backend::{Backend, QueryResult, QueryStatus, SolveContext}; + +/// Runs several simplifier backends in sequence, feeding each stage the +/// previous stage's output expression. A stage that declines (returns +/// `UNKNOWN`, errors, or produces an invalid artifact) is skipped and the +/// pipeline continues from the last good result. +/// +/// This lets complementary simplifiers compose: Rumba's fast linear-MBA +/// rewriting first, CoBRA's deeper worklist pipeline on what remains. +#[derive(Clone)] +pub struct SimplifyChainBackend { + stages: Vec>, +} + +impl SimplifyChainBackend { + pub fn new(stages: Vec>) -> Self { + Self { stages } + } + + pub fn stages(&self) -> &[Arc] { + &self.stages + } + + fn run( + &self, + request: &BinaryRequest, + context: Option<&SolveContext>, + ) -> smt_wire::Result { + let target = request + .target_ref() + .ok_or_else(|| WireError::invalid("simplify request", "missing target_node"))?; + let mut current = SimplifyBlock { + expression: request.expression.clone(), + target_node: target, + }; + for stage in &self.stages { + if context.is_some_and(SolveContext::is_cancelled) { + return Ok(cancelled_result()); + } + let Ok(stage_request) = request_for_block(request, ¤t) else { + break; + }; + let result = match context { + Some(context) => stage.handle_with_context(&stage_request, context), + None => stage.handle(&stage_request), + }; + let Ok(result) = result else { + continue; + }; + if result.status != QueryStatus::Simplified + || result.validate_artifacts_for(&stage_request).is_err() + { + continue; + } + if let Some(block) = result.simplify { + current = block; + } + } + // Match the backend cancellation convention: a cancelled request + // returns an inconclusive result, never a (vacuously valid) identity + // simplification that a racing layer could adopt. Cancellation can + // also arrive while the final stage runs, so check again on the way + // out, not only before each stage. + if context.is_some_and(SolveContext::is_cancelled) { + return Ok(cancelled_result()); + } + Ok(QueryResult::simplified(current)) + } +} + +fn cancelled_result() -> QueryResult { + QueryResult::unknown("simplify chain request cancelled before completion") +} + +/// Rewrap a stage's output as the next stage's input request, keeping the +/// original envelope but dropping assertions: SIMPLIFY operates on the target +/// expression alone. +fn request_for_block( + request: &BinaryRequest, + block: &SimplifyBlock, +) -> smt_wire::Result { + BinaryRequest::new( + request.envelope.request_id, + Command::Simplify, + request.envelope.flags, + request.envelope.budget_ms, + block.expression.clone(), + Vec::new(), + Vec::new(), + Vec::new(), + Some(block.target_node), + ) +} + +impl Backend for SimplifyChainBackend { + fn name(&self) -> &'static str { + "simplify-chain" + } + + fn handle(&self, request: &BinaryRequest) -> smt_wire::Result { + match request.envelope.command { + Command::Simplify => self.run(request, None), + Command::Solve | Command::Minimize | Command::Maximize => Ok(QueryResult::unknown( + "simplify-chain only supports SIMPLIFY", + )), + } + } + + fn handle_with_context( + &self, + request: &BinaryRequest, + context: &SolveContext, + ) -> smt_wire::Result { + match request.envelope.command { + Command::Simplify => self.run(request, Some(context)), + Command::Solve | Command::Minimize | Command::Maximize => Ok(QueryResult::unknown( + "simplify-chain only supports SIMPLIFY", + )), + } + } +} diff --git a/crates/smt-server/tests/cobra_simplify.rs b/crates/smt-server/tests/cobra_simplify.rs new file mode 100644 index 0000000..9482d70 --- /dev/null +++ b/crates/smt-server/tests/cobra_simplify.rs @@ -0,0 +1,295 @@ +//! CoBRA simplifies Lean-certified MBA islands, leaves everything else +//! unchanged, and composes with Rumba through the simplify chain. + +use std::sync::Arc; + +use smt_server::{handle_binary_frame, Backend, CobraBackend, RumbaBackend, SimplifyChainBackend}; +use smt_wire::raw::{tag, BinaryResponse, ExprBuilder, NodeRef, SimplifyBlock, Status}; + +/// Count occurrences of each tag in the simplified result expression. +fn tag_counts(payload: &[u8]) -> Vec { + let block = SimplifyBlock::decode(payload).unwrap(); + let buffer = block.expression_buffer().unwrap(); + let view = buffer.view().unwrap(); + let mut tags = Vec::new(); + for index in 0..view.node_count() { + tags.push(view.node(index).unwrap().tag); + } + tags +} + +fn simplify_with(backend: &dyn Backend, builder: &ExprBuilder, target: NodeRef) -> Vec { + let request = builder.build_simplify_request(1, target).unwrap(); + let response = BinaryResponse::parse( + &handle_binary_frame(&request, backend) + .unwrap() + .encode() + .unwrap(), + ) + .unwrap(); + assert_eq!(response.envelope.status, Status::Simplified); + response.payload +} + +fn simplify(builder: &ExprBuilder, target: NodeRef) -> Vec { + simplify_with(&CobraBackend::default(), builder, target) +} + +/// `(x ^ y) + 2 * (x & y)` -> `x + y`, backed by CoBRA's Lean certificate. +#[test] +fn simplifies_certified_mba_identity() { + let mut b = ExprBuilder::new(); + let x = b.bv_var("x", 64).unwrap(); + let y = b.bv_var("y", 64).unwrap(); + let two = b.bv_const(2, 64).unwrap(); + let xor = b.bv_xor(x, y).unwrap(); + let and = b.bv_and(x, y).unwrap(); + let mul = b.bv_mul(two, and).unwrap(); + let target = b.bv_add(xor, mul).unwrap(); + + let tags = tag_counts(&simplify(&b, target)); + let count = |t: u8| tags.iter().filter(|&&n| n == t).count(); + assert_eq!(count(tag::BV_XOR), 0, "xor not eliminated"); + assert_eq!(count(tag::BV_AND), 0, "and not eliminated"); + assert_eq!(count(tag::BV_MUL), 0, "mul not eliminated"); + assert!(count(tag::BV_ADD) >= 1, "reduced x+y add missing"); +} + +/// The same island simplifies at a narrow width; CoBRA replays its +/// certificates per bit-width. +#[test] +fn simplifies_mba_identity_at_narrow_width() { + let mut b = ExprBuilder::new(); + let x = b.bv_var("x", 16).unwrap(); + let y = b.bv_var("y", 16).unwrap(); + let or = b.bv_or(x, y).unwrap(); + let and = b.bv_and(x, y).unwrap(); + let target = b.bv_add(or, and).unwrap(); + + let tags = tag_counts(&simplify(&b, target)); + let count = |t: u8| tags.iter().filter(|&&n| n == t).count(); + assert_eq!(count(tag::BV_OR), 0, "or not eliminated"); + assert_eq!(count(tag::BV_AND), 0, "and not eliminated"); + assert!(count(tag::BV_ADD) >= 1, "reduced x+y add missing"); +} + +/// `concat(0xFF, (x^y) + 2*(x&y))` -> the MBA arm must collapse while the +/// `concat` skeleton is preserved. +#[test] +fn simplifies_mba_island_nested_under_concat() { + let mut b = ExprBuilder::new(); + let x = b.bv_var("x", 64).unwrap(); + let y = b.bv_var("y", 64).unwrap(); + let two = b.bv_const(2, 64).unwrap(); + let xor = b.bv_xor(x, y).unwrap(); + let and = b.bv_and(x, y).unwrap(); + let mul = b.bv_mul(two, and).unwrap(); + let mba = b.bv_add(xor, mul).unwrap(); + let hi = b.bv_const(0xFF, 16).unwrap(); + let target = b.bv_concat(hi, mba).unwrap(); + + let tags = tag_counts(&simplify(&b, target)); + let count = |t: u8| tags.iter().filter(|&&n| n == t).count(); + assert_eq!(count(tag::BV_XOR), 0, "xor not eliminated"); + assert_eq!(count(tag::BV_AND), 0, "and not eliminated"); + assert_eq!(count(tag::BV_MUL), 0, "mul not eliminated"); + assert_eq!(count(tag::BV_CONCAT), 1, "concat skeleton not preserved"); +} + +/// A non-MBA subtree becomes an opaque boundary; the identity `boundary ^ 0` +/// island around it must still collapse without touching the `udiv` inside. +#[test] +fn preserves_non_mba_boundary_subtrees() { + let mut b = ExprBuilder::new(); + let x = b.bv_var("x", 64).unwrap(); + let y = b.bv_var("y", 64).unwrap(); + let div = b.bv_udiv(x, y).unwrap(); + let or = b.bv_or(div, div).unwrap(); + let target = b.bv_or(or, div).unwrap(); + + let tags = tag_counts(&simplify(&b, target)); + let count = |t: u8| tags.iter().filter(|&&n| n == t).count(); + assert_eq!(count(tag::BV_UDIV), 1, "udiv boundary not preserved"); + assert_eq!(count(tag::BV_OR), 0, "or self-absorption not eliminated"); +} + +/// A wire variable that collides with the synthetic boundary-name prefix must +/// not break variable naming: the result is still a valid simplification. +#[test] +fn tolerates_boundary_name_collisions() { + let mut b = ExprBuilder::new(); + let x = b.bv_var("x", 64).unwrap(); + let y = b.bv_var("y", 64).unwrap(); + let clash = b.bv_var("__cobra_boundary0", 64).unwrap(); + let div = b.bv_udiv(x, y).unwrap(); + let add = b.bv_add(div, clash).unwrap(); + let target = b.bv_xor(add, add).unwrap(); + + let tags = tag_counts(&simplify(&b, target)); + let count = |t: u8| tags.iter().filter(|&&n| n == t).count(); + assert_eq!( + count(tag::BV_XOR), + 0, + "xor self-cancellation not eliminated" + ); +} + +/// Islands CoBRA cannot improve keep their structural skeleton intact. +#[test] +fn unimproved_islands_keep_their_skeleton() { + let mut b = ExprBuilder::new(); + let x = b.bv_var("x", 64).unwrap(); + let mask = b.bv_const(0xFF, 64).unwrap(); + let and = b.bv_and(x, mask).unwrap(); + let shifted = b.bv_shl(and, mask).unwrap(); + let target = b.bv_add(shifted, and).unwrap(); + + let tags = tag_counts(&simplify(&b, target)); + let count = |t: u8| tags.iter().filter(|&&n| n == t).count(); + assert_eq!(count(tag::BV_SHL), 1, "shl skeleton not preserved"); + assert_eq!(count(tag::BV_ADD), 1, "add skeleton not preserved"); +} + +/// A deep doubling chain shares each node on both sides of an add, so tree +/// conversion without the shared island budget would take ~2^48 visits. The +/// budget must cut the attempts and fall back to a structural copy quickly. +#[test] +fn bounds_conversion_work_on_shared_dags() { + let mut b = ExprBuilder::new(); + let x = b.bv_var("x", 64).unwrap(); + let mut acc = x; + for _ in 0..48 { + acc = b.bv_add(acc, acc).unwrap(); + } + + let tags = tag_counts(&simplify(&b, acc)); + let count = |t: u8| tags.iter().filter(|&&n| n == t).count(); + assert_eq!(count(tag::BV_ADD), 48, "doubling chain not preserved"); + assert_eq!(count(tag::BV_VAR), 1, "shared variable not preserved"); +} + +/// The chain feeds Rumba's output into CoBRA; each backend's specialty lands +/// in one pass through the chain. +#[test] +fn chain_composes_rumba_and_cobra() { + let chain = SimplifyChainBackend::new(vec![ + Arc::new(RumbaBackend), + Arc::new(CobraBackend::default()), + ]); + + let mut b = ExprBuilder::new(); + let x = b.bv_var("x", 64).unwrap(); + let y = b.bv_var("y", 64).unwrap(); + let two = b.bv_const(2, 64).unwrap(); + let xor = b.bv_xor(x, y).unwrap(); + let and = b.bv_and(x, y).unwrap(); + let mul = b.bv_mul(two, and).unwrap(); + let target = b.bv_add(xor, mul).unwrap(); + + let tags = tag_counts(&simplify_with(&chain, &b, target)); + let count = |t: u8| tags.iter().filter(|&&n| n == t).count(); + assert_eq!(count(tag::BV_XOR), 0, "xor not eliminated"); + assert_eq!(count(tag::BV_AND), 0, "and not eliminated"); + assert_eq!(count(tag::BV_MUL), 0, "mul not eliminated"); + assert!(count(tag::BV_ADD) >= 1, "reduced x+y add missing"); +} + +/// A chain still produces a valid identity result when every stage declines. +#[test] +fn chain_falls_back_to_identity() { + let chain = SimplifyChainBackend::new(vec![Arc::new(CobraBackend::default())]); + + let mut b = ExprBuilder::new(); + let x = b.bv_var("x", 64).unwrap(); + let y = b.bv_var("y", 64).unwrap(); + let target = b.bv_udiv(x, y).unwrap(); + + let tags = tag_counts(&simplify_with(&chain, &b, target)); + let count = |t: u8| tags.iter().filter(|&&n| n == t).count(); + assert_eq!(count(tag::BV_UDIV), 1, "udiv not preserved"); +} + +/// Cancellation arriving while the final stage runs must also surface as an +/// inconclusive result, not as the last good block. +#[test] +fn chain_observes_mid_stage_cancellation() { + use smt_server::{CancellationToken, QueryResult, SolveContext}; + use smt_wire::raw::BinaryRequest; + + /// Stage that cancels the request while handling it, like a racing layer + /// reacting to another backend's win mid-flight. + struct CancelDuringHandle(CancellationToken); + + impl Backend for CancelDuringHandle { + fn name(&self) -> &'static str { + "cancel-during-handle" + } + + fn handle(&self, _request: &BinaryRequest) -> smt_wire::Result { + Ok(QueryResult::unknown("cancelled mid-stage")) + } + + fn handle_with_context( + &self, + request: &BinaryRequest, + _context: &SolveContext, + ) -> smt_wire::Result { + self.0.cancel(); + self.handle(request) + } + } + + let cancellation = CancellationToken::new(); + let chain = SimplifyChainBackend::new(vec![Arc::new(CancelDuringHandle(cancellation.clone()))]); + + let mut b = ExprBuilder::new(); + let x = b.bv_var("x", 64).unwrap(); + let y = b.bv_var("y", 64).unwrap(); + let target = b.bv_xor(x, y).unwrap(); + let request = BinaryRequest::parse(&b.build_simplify_request(1, target).unwrap()).unwrap(); + + let context = SolveContext::new(cancellation); + let result = chain.handle_with_context(&request, &context).unwrap(); + assert!(!result.is_conclusive(), "chain returned {result:?}"); + assert!( + result + .message + .as_deref() + .unwrap_or_default() + .contains("cancelled"), + "missing cancellation message: {result:?}" + ); +} + +/// A pre-cancelled context yields an inconclusive result, matching the +/// backend cancellation convention, instead of an identity simplification. +#[test] +fn chain_observes_pre_cancelled_context() { + use smt_server::{CancellationToken, SolveContext}; + use smt_wire::raw::BinaryRequest; + + let chain = SimplifyChainBackend::new(vec![ + Arc::new(RumbaBackend), + Arc::new(CobraBackend::default()), + ]); + + let mut b = ExprBuilder::new(); + let x = b.bv_var("x", 64).unwrap(); + let y = b.bv_var("y", 64).unwrap(); + let target = b.bv_xor(x, y).unwrap(); + let request = BinaryRequest::parse(&b.build_simplify_request(1, target).unwrap()).unwrap(); + + let cancellation = CancellationToken::new(); + cancellation.cancel(); + let context = SolveContext::new(cancellation); + let result = chain.handle_with_context(&request, &context).unwrap(); + assert!(!result.is_conclusive(), "chain returned {result:?}"); + assert!( + result + .message + .as_deref() + .unwrap_or_default() + .contains("cancelled"), + "missing cancellation message: {result:?}" + ); +} diff --git a/docs/architecture.md b/docs/architecture.md index b081130..38c2b35 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -13,7 +13,7 @@ The server is intentionally stateless at the protocol boundary. Every request co | Path | Role | |---|---| | `crates/smt-wire` | Rust protocol/client crate. It owns wire constants, expression validation, request/response codecs, high-level Rust `Context`/term APIs, and a TCP client. | -| `crates/smt-server` | TCP server, binary/text dispatch, response cache, backend trait, backend racing, and integrations for Z3, binbit, qfbvsmtrs, and Rumba. | +| `crates/smt-server` | TCP server, binary/text dispatch, response cache, backend trait, backend racing, and integrations for Z3, binbit, qfbvsmtrs, Rumba, and CoBRA. | | `crates/smt-qfbv-smtlib` | Shared solver-agnostic QF_BV/Bool SMT-LIB frontend. It parses with `yaspar` and lowers through a sink trait. | | `crates/qfbvsmtrs` | Standalone pure-Rust QF_BV solver. It can be used directly, through its SMT-LIB frontend, or as a server backend through the `smt-wire` bridge. | | `python` | Dependency-free Python client package (`smt_wire.py`) plus examples/tests. Install from this repo with pip's `#subdirectory=python` support. | @@ -73,7 +73,9 @@ The shipped binary in `crates/smt-server/src/main.rs` builds this backend stack: ```text CommandRouterBackend -├── SIMPLIFY -> RumbaBackend +├── SIMPLIFY -> SimplifyChainBackend +│ ├── RumbaBackend +│ └── CobraBackend └── SOLVE / MINIMIZE / MAXIMIZE -> RacingBackend(default budget: 30s) ├── Z3Backend @@ -87,6 +89,8 @@ Backend responsibilities: - `BinbitBackend` translates validated wire IR to `binbit` and supports solve, model extraction, named unsat cores, and optimization helpers. - `QfbvsmtrsBackend` lowers wire requests into the standalone qfbvsmtrs IR and uses the pure-Rust bit-blast/SAT pipeline. - `RumbaBackend` handles `SIMPLIFY` for supported 64-bit-or-smaller MBA expression islands. Unsupported simplifications return the original target expression rather than a wrong rewrite. +- `CobraBackend` handles `SIMPLIFY` with the [CoBRA](https://github.com/binsnake/cobra) worklist-driven MBA simplifier, using the same island extraction as Rumba. By default it only adopts rewrites CoBRA backs with a replayable Lean certificate; a spot-checked mode trades that guarantee for a higher simplification rate. +- `SimplifyChainBackend` runs simplifier backends in sequence, feeding each stage the previous stage's output and skipping stages that decline. `RacingBackend` returns the first conclusive answer and logs later disagreements for investigation. `UNKNOWN` is safe and means no backend produced a conclusive answer within the applicable budget.