patronus 0.38.1

// Copyright 2023-2024 The Regents of the University of California
// Copyright 2024 Cornell University
// Copyright 2026 The Hong Kong University of Science and Technology
// released under BSD 3-Clause License
// author: Kevin Laeufer <laeufer@cornell.edu>
// author: Guangyu Hu <ghuae@connect.ust.hk>, The Hong Kong University of Science and Technology

use super::parse::{
    DEFAULT_BAD_STATE_PREFIX, DEFAULT_CONSTRAINT_PREFIX, PARSER_RESERVED_NAMES, unique_name,
};
use crate::VERSION;
use crate::expr::*;
use crate::system::{State, TransitionSystem};
use baa::BitVecOps;
use regex::Regex;
use rustc_hash::{FxHashMap, FxHashSet};
use std::io::Write;

/// The patronus parser assigns default names of the form `_state`, `_state_0`, `_state_1`,
/// ... to symbols that arrive without an explicit name in the source BTOR2. Re-serializing
/// those synthesized names and re-parsing them pollutes the parser's unique-name namespace
/// and forces unwanted `_N` suffixes on unrelated symbols, which breaks round-trips. We
/// filter such names back out on the way to BTOR2 and let the parser regenerate them.
fn is_autogen_name(name: &str) -> bool {
    AUTOGEN_NAME_REGEX.is_match(name)
}

lazy_static! {
    static ref AUTOGEN_NAME_REGEX: Regex = {
        let prefix = PARSER_RESERVED_NAMES
            .iter()
            .map(|n| format!("({n})"))
            .collect::<Vec<_>>()
            .join("|");
        Regex::new(&format!("^({prefix})(_\\d+)?$")).unwrap()
    };
}

pub fn serialize(
    ctx: &Context,
    writer: &mut impl Write,
    sys: &TransitionSystem,
) -> std::io::Result<()> {
    Serializer::new(ctx, writer).serialize_sys(sys)
}

pub fn serialize_to_str(ctx: &Context, sys: &TransitionSystem) -> String {
    let mut buf = Vec::new();
    serialize(ctx, &mut buf, sys).expect("Failed to write to string!");
    String::from_utf8(buf).expect("Failed to read string we wrote!")
}

struct Serializer<'a, W: Write> {
    ctx: &'a Context,
    writer: &'a mut W,
    next_id: u64,
    sort_ids: FxHashMap<Type, u64>,
    expr_ids: SparseExprMap<Option<u64>>,
    /// precomputed output/bad/constraint labels
    label_names: FxHashSet<String>,
    /// Expressions whose canonical name is restored via a trailing
    /// `uext <sort> <e> 0 <name>` alias line; they are emitted unnamed inline.
    alias_needed: FxHashSet<ExprRef>,
}

struct LabelNames {
    outputs: Vec<String>,
    constraints: Vec<String>,
    bads: Vec<String>,
}

impl LabelNames {
    fn all(&self) -> FxHashSet<String> {
        self.outputs
            .iter()
            .chain(self.constraints.iter())
            .chain(self.bads.iter())
            .cloned()
            .collect()
    }
}

impl<'a, W: Write> Serializer<'a, W> {
    fn new(ctx: &'a Context, writer: &'a mut W) -> Self {
        Serializer {
            ctx,
            writer,
            next_id: 1,
            sort_ids: FxHashMap::default(),
            expr_ids: SparseExprMap::default(),
            label_names: FxHashSet::default(),
            alias_needed: FxHashSet::default(),
        }
    }

    fn new_id(&mut self) -> u64 {
        let id = self.next_id;
        self.next_id += 1;
        id
    }

    fn sort_id(&mut self, tpe: Type) -> std::io::Result<u64> {
        if let Some(&id) = self.sort_ids.get(&tpe) {
            return Ok(id);
        }
        let id = match tpe {
            Type::BV(w) => {
                let id = self.new_id();
                writeln!(self.writer, "{id} sort bitvec {w}")?;
                id
            }
            Type::Array(ArrayType {
                index_width,
                data_width,
            }) => {
                let ix = self.sort_id(Type::BV(index_width))?;
                let dx = self.sort_id(Type::BV(data_width))?;
                let id = self.new_id();
                writeln!(self.writer, "{id} sort array {ix} {dx}")?;
                id
            }
        };
        self.sort_ids.insert(tpe, id);
        Ok(id)
    }

    fn serialize_sys(&mut self, sys: &TransitionSystem) -> std::io::Result<()> {
        writeln!(
            self.writer,
            "; btor2 description of `{}` generated by patronus {}",
            sys.name,
            VERSION.unwrap_or_default()
        )?;

        // precompute label names for stable round-trips
        let labels = compute_label_names(self.ctx, sys);
        self.label_names = labels.all();

        // Labels (outputs/bads/constraints) directly referencing an expression whose
        // canonical name differs from the label's name. For each such expression we
        // defer the name to a trailing `uext <e> 0 <name>` alias line: otherwise the
        // label's `sys.names[e]` overwrite (plus `improve_state_names`, for states)
        // would lose the original name on re-parse.
        self.alias_needed = compute_alias_needed(self.ctx, sys, &labels);

        // 1. declare inputs
        for &input in sys.inputs.iter() {
            let tpe = input.get_type(self.ctx);
            let sort = self.sort_id(tpe)?;
            let id = self.new_id();
            let raw = self.ctx[input].get_symbol_name(self.ctx).unwrap_or("");
            let name = decl_name(raw, &self.label_names);
            writeln!(self.writer, "{id} input {sort}{}", name_suffix(name))?;
            self.expr_ids[input] = Some(id);
        }

        // 2-4. For each state in declaration order, interleave:
        //    (a) the state's init expression tree (before the state decl, as required
        //        by BTOR2: the init line's expr operand must have a smaller id than the
        //        state id itself),
        //    (b) the state decl,
        //    (c) the corresponding `init` line.
        //
        //    Emitting one state at a time (rather than pre-walking all inits first) is
        //    required to handle inputs like `const_array_example.btor`, where one state's
        //    init references an earlier state (e.g. `mem_n` initialized to `write(mem,
        //    addr, data)`). By the time we process `mem_n`, `mem` is already declared,
        //    so its symbol resolves and its state id is smaller than `mem_n`'s.
        let mut state_ids: Vec<u64> = Vec::with_capacity(sys.states.len());
        for state in sys.states.iter() {
            let state_tpe = state.symbol.get_type(self.ctx);
            let state_sort = self.sort_id(state_tpe)?;

            // (a) init expression tree, if any.
            let init_id = if let Some(init) = state.init {
                Some(self.emit_state_init(sys, state, init)?)
            } else {
                None
            };

            // (b) state decl.
            let state_id = self.new_id();
            let raw = self.ctx[state.symbol]
                .get_symbol_name(self.ctx)
                .unwrap_or("");
            let name = if self.alias_needed.contains(&state.symbol) {
                "" // trailing alias line will reassert the name
            } else {
                decl_name(raw, &self.label_names)
            };
            writeln!(
                self.writer,
                "{state_id} state {state_sort}{}",
                name_suffix(name)
            )?;
            self.expr_ids[state.symbol] = Some(state_id);
            state_ids.push(state_id);

            // (c) init line. By construction, init_id < state_id.
            if let Some(init_id) = init_id {
                let line_id = self.new_id();
                writeln!(
                    self.writer,
                    "{line_id} init {state_sort} {state_id} {init_id}"
                )?;
            }
        }

        // 5. outputs
        for (out, name) in sys.outputs.iter().zip(labels.outputs.iter()) {
            let body_id = self.emit_expr(sys, out.expr)?;
            let id = self.new_id();
            writeln!(self.writer, "{id} output {body_id}{}", name_suffix(name))?;
        }

        // 6. constraints
        for (&c, name) in sys.constraints.iter().zip(labels.constraints.iter()) {
            let body_id = self.emit_expr(sys, c)?;
            let id = self.new_id();
            writeln!(
                self.writer,
                "{id} constraint {body_id}{}",
                name_suffix(name)
            )?;
        }

        // 7. bad states
        for (&b, name) in sys.bad_states.iter().zip(labels.bads.iter()) {
            let body_id = self.emit_expr(sys, b)?;
            let id = self.new_id();
            writeln!(self.writer, "{id} bad {body_id}{}", name_suffix(name))?;
        }

        // 8. trailing aliases: reassert the canonical name of every flagged symbol
        //    (state *or* intermediate expression). `uext ... 0 <name>` is the idiomatic
        //    pattern the parser already recognizes — for states it feeds into
        //    `improve_state_names`; for non-state expressions it just overwrites
        //    `sys.names[e]` with the intended debug name, restoring what the preceding
        //    output/bad/constraint line overwrote. Because the flagged node was emitted
        //    without its name (see `emit_expr` / state decl above), `name` is guaranteed
        //    unique in the parser's unique-names set at this point.
        // Collect first so we don't borrow `self.alias_needed` while mutably borrowing
        // `self.writer`/`self.sort_ids` inside the loop. Sort by the already-assigned
        // BTOR2 id so alias ordering is deterministic across runs (a HashSet's iter
        // order is not, which would cause spurious round-trip diffs).
        let mut alias_targets: Vec<(ExprRef, u64)> = self
            .alias_needed
            .iter()
            .filter_map(|&e| self.expr_ids[e].map(|id| (e, id)))
            .collect();
        alias_targets.sort_by_key(|&(_, id)| id);
        for (e, id_of_target) in alias_targets {
            let name = expr_canonical_name(self.ctx, sys, e);
            if name.is_empty() {
                continue;
            }
            let sort = self.sort_id(e.get_type(self.ctx))?;
            let line_id = self.new_id();
            writeln!(self.writer, "{line_id} uext {sort} {id_of_target} 0 {name}")?;
        }

        // 9. next expressions and lines (next may reference states, so it has to come
        //    after state decls).
        for (state, &state_id) in sys.states.iter().zip(state_ids.iter()) {
            if let Some(next) = state.next {
                let state_sort = self.sort_id(state.symbol.get_type(self.ctx))?;
                let next_id = self.emit_expr(sys, next)?;
                let id = self.new_id();
                writeln!(self.writer, "{id} next {state_sort} {state_id} {next_id}")?;
            }
        }

        Ok(())
    }

    /// BTOR2 lets you initialize an array state from a bit-vector value directly: the
    /// `init` line carries the array sort but a BV-sort node. The parser re-wraps such
    /// cases into `ArrayConstant`, so on the way out we unwrap the BV and let the init
    /// line handle the rest.
    fn emit_state_init(
        &mut self,
        sys: &TransitionSystem,
        state: &State,
        init: ExprRef,
    ) -> std::io::Result<u64> {
        if state.symbol.get_type(self.ctx).is_array()
            && let Expr::ArrayConstant { e, .. } = self.ctx[init]
        {
            return self.emit_expr(sys, e);
        }
        self.emit_expr(sys, init)
    }

    /// Emit this expression (and all sub-expressions) in post-order. Symbols must already
    /// be registered via `expr_ids` (inputs/states).
    fn emit_expr(&mut self, sys: &TransitionSystem, e: ExprRef) -> std::io::Result<u64> {
        if let Some(id) = self.expr_ids[e] {
            return Ok(id);
        }

        let expr = self.ctx[e].clone();

        if expr.is_symbol() {
            let name = expr.get_symbol_name(self.ctx).unwrap_or("?");
            panic!(
                "encountered unregistered symbol `{name}` while serializing btor2 — \
                 every symbol must be declared as an input or state"
            );
        }

        // emit children first so they have assigned ids
        let mut children: Vec<ExprRef> = Vec::with_capacity(3);
        expr.for_each_child(|c| children.push(*c));
        let child_ids: Vec<u64> = children
            .into_iter()
            .map(|c| self.emit_expr(sys, c))
            .collect::<Result<_, _>>()?;

        let tpe = expr.get_type(self.ctx);
        let sort = self.sort_id(tpe)?;
        let id = self.new_id();

        // Preserve yosys-style debug names carried on intermediate expressions
        // (e.g. `uext 1 11 0 axis_reg_inst.s_axis_tready`). These names sit in
        // `sys.names` after parsing. We omit the name in two cases:
        //   * It collides with a label name (output/bad/constraint). Emitting it
        //     would pollute the parser's unique-name set and force the label to
        //     get an `_N` suffix on re-parse.
        //   * The expression is in `alias_needed`: the canonical name is going
        //     to be re-asserted by a trailing `uext 0` alias line anyway.
        let tail = if self.alias_needed.contains(&e) {
            String::new()
        } else {
            sys.names[e]
                .map(|sr| &self.ctx[sr])
                .filter(|name| !self.label_names.contains(name.as_str()))
                .map(|s| name_suffix(s))
                .unwrap_or_default()
        };

        write_node(self.writer, self.ctx, id, sort, &expr, &child_ids, &tail)?;

        self.expr_ids[e] = Some(id);
        Ok(id)
    }
}

/// Pads name with a space for use at the end of a btor2 line
fn name_suffix(name: &str) -> String {
    if name.is_empty() {
        String::new()
    } else {
        format!(" {name}")
    }
}

/// Pick the name to emit on a symbol's declaration line. We omit the name in two cases:
/// (1) the symbol carries an autogenerated default prefix (the parser will regenerate it),
/// (2) the name would clash with an output/bad/constraint label (in which case the parser's
/// `improve_state_names` pass will propagate the name back onto the symbol on re-parse).
fn decl_name<'a>(raw: &'a str, label_names: &FxHashSet<String>) -> &'a str {
    if raw.is_empty() || is_autogen_name(raw) || label_names.contains(raw) {
        ""
    } else {
        raw
    }
}

/// Returns the set of expressions that need a trailing `uext 0` alias line to preserve
/// their name across a round-trip.
///
/// An expression `e` is flagged when it is *directly* referenced by some
/// output/bad/constraint whose label name differs from the name carried by `e` (either
/// the symbol name for a state/input, or the yosys-style debug name in `sys.names[e]`
/// for an intermediate expression). Without the alias:
///   * For a state, the parser's `improve_state_names` pass would override the state's
///     name with that label's name on re-parse.
///   * For any other expression, parser's overwrite of `sys.names[e]` from the
///     label line would simply replace the debug name, and on re-serialization we'd
///     drop it under the label-collision filter.
fn compute_alias_needed(
    ctx: &Context,
    sys: &TransitionSystem,
    labels: &LabelNames,
) -> FxHashSet<ExprRef> {
    // Track the *last* label (in emit order: outputs, then constraints, then bads)
    // that directly references each expression. The parser's `sys.names[e]` after
    // re-parse ends up equal to that last label, so we only need a trailing alias
    // when the canonical name differs from it.
    let mut last_label: FxHashMap<ExprRef, String> = FxHashMap::default();

    for (o, label) in sys.outputs.iter().zip(labels.outputs.iter()) {
        last_label.insert(o.expr, label.clone());
    }
    for (&e, label) in sys.constraints.iter().zip(labels.constraints.iter()) {
        last_label.insert(e, label.clone());
    }
    for (&e, label) in sys.bad_states.iter().zip(labels.bads.iter()) {
        last_label.insert(e, label.clone());
    }

    let mut out = FxHashSet::default();
    for (e, label) in last_label {
        let name = expr_canonical_name(ctx, sys, e);
        if name.is_empty() || is_autogen_name(&name) {
            continue;
        }
        if name != label {
            out.insert(e);
        }
    }
    out
}

/// The "canonical" name for an expression that should survive a round-trip. For a
/// symbol (state/input), this is the symbol's own name. For any other expression,
/// this is the yosys-style debug label stashed in `sys.names[e]` by the parser.
fn expr_canonical_name(ctx: &Context, sys: &TransitionSystem, e: ExprRef) -> String {
    if let Some(name) = ctx[e].get_symbol_name(ctx) {
        return name.to_string();
    }
    sys.names[e].map(|s| ctx[s].clone()).unwrap_or_default()
}

fn compute_label_names(ctx: &Context, sys: &TransitionSystem) -> LabelNames {
    let mut used = PARSER_RESERVED_NAMES
        .iter()
        .map(|s| s.to_string())
        .collect();

    let outputs = sys
        .outputs
        .iter()
        .map(|out| unique_name(&ctx[out.name], &mut used))
        .collect();
    let constraints = sys
        .constraints
        .iter()
        .map(|&e| {
            unique_name(
                &label_name_base(ctx, sys, e, DEFAULT_CONSTRAINT_PREFIX),
                &mut used,
            )
        })
        .collect();
    let bads = sys
        .bad_states
        .iter()
        .map(|&e| {
            unique_name(
                &label_name_base(ctx, sys, e, DEFAULT_BAD_STATE_PREFIX),
                &mut used,
            )
        })
        .collect();

    LabelNames {
        outputs,
        constraints,
        bads,
    }
}

fn label_name_base(ctx: &Context, sys: &TransitionSystem, e: ExprRef, default: &str) -> String {
    let name = ctx[e]
        .get_symbol_name(ctx)
        .map(|s| s.to_string())
        .or_else(|| sys.names[e].map(|s| ctx[s].clone()));
    match name {
        Some(name) if !is_autogen_name(&name) => name,
        _ => default.to_string(),
    }
}

fn write_node<W: Write>(
    writer: &mut W,
    ctx: &Context,
    id: u64,
    sort: u64,
    expr: &Expr,
    children: &[u64],
    tail: &str,
) -> std::io::Result<()> {
    match expr {
        Expr::BVSymbol { .. } | Expr::ArraySymbol { .. } => {
            unreachable!("symbols are handled by the caller and never emitted as regular nodes")
        }

        Expr::BVLiteral(value) => write_bv_literal(writer, ctx, id, sort, *value)?,

        Expr::BVZeroExt { by, .. } => write!(writer, "{id} uext {sort} {} {by}", children[0])?,
        Expr::BVSignExt { by, .. } => write!(writer, "{id} sext {sort} {} {by}", children[0])?,
        Expr::BVSlice { hi, lo, .. } => {
            write!(writer, "{id} slice {sort} {} {hi} {lo}", children[0])?
        }

        Expr::BVNot(_, _) => write!(writer, "{id} not {sort} {}", children[0])?,
        Expr::BVNegate(_, _) => write!(writer, "{id} neg {sort} {}", children[0])?,

        Expr::BVEqual(_, _) | Expr::ArrayEqual(_, _) => {
            write!(writer, "{id} eq {sort} {} {}", children[0], children[1])?
        }
        Expr::BVImplies(_, _) => write!(
            writer,
            "{id} implies {sort} {} {}",
            children[0], children[1]
        )?,
        Expr::BVGreater(_, _) => write!(writer, "{id} ugt {sort} {} {}", children[0], children[1])?,
        Expr::BVGreaterSigned(_, _, _) => {
            write!(writer, "{id} sgt {sort} {} {}", children[0], children[1])?
        }
        Expr::BVGreaterEqual(_, _) => {
            write!(writer, "{id} ugte {sort} {} {}", children[0], children[1])?
        }
        Expr::BVGreaterEqualSigned(_, _, _) => {
            write!(writer, "{id} sgte {sort} {} {}", children[0], children[1])?
        }
        Expr::BVConcat(_, _, _) => {
            write!(writer, "{id} concat {sort} {} {}", children[0], children[1])?
        }

        Expr::BVAnd(_, _, _) => write!(writer, "{id} and {sort} {} {}", children[0], children[1])?,
        Expr::BVOr(_, _, _) => write!(writer, "{id} or {sort} {} {}", children[0], children[1])?,
        Expr::BVXor(_, _, _) => write!(writer, "{id} xor {sort} {} {}", children[0], children[1])?,
        Expr::BVShiftLeft(_, _, _) => {
            write!(writer, "{id} sll {sort} {} {}", children[0], children[1])?
        }
        Expr::BVArithmeticShiftRight(_, _, _) => {
            write!(writer, "{id} sra {sort} {} {}", children[0], children[1])?
        }
        Expr::BVShiftRight(_, _, _) => {
            write!(writer, "{id} srl {sort} {} {}", children[0], children[1])?
        }
        Expr::BVAdd(_, _, _) => write!(writer, "{id} add {sort} {} {}", children[0], children[1])?,
        Expr::BVMul(_, _, _) => write!(writer, "{id} mul {sort} {} {}", children[0], children[1])?,
        Expr::BVSignedDiv(_, _, _) => {
            write!(writer, "{id} sdiv {sort} {} {}", children[0], children[1])?
        }
        Expr::BVUnsignedDiv(_, _, _) => {
            write!(writer, "{id} udiv {sort} {} {}", children[0], children[1])?
        }
        Expr::BVSignedMod(_, _, _) => {
            write!(writer, "{id} smod {sort} {} {}", children[0], children[1])?
        }
        Expr::BVSignedRem(_, _, _) => {
            write!(writer, "{id} srem {sort} {} {}", children[0], children[1])?
        }
        Expr::BVUnsignedRem(_, _, _) => {
            write!(writer, "{id} urem {sort} {} {}", children[0], children[1])?
        }
        Expr::BVSub(_, _, _) => write!(writer, "{id} sub {sort} {} {}", children[0], children[1])?,

        Expr::BVArrayRead { .. } => {
            write!(writer, "{id} read {sort} {} {}", children[0], children[1])?
        }

        Expr::BVIte { .. } | Expr::ArrayIte { .. } => write!(
            writer,
            "{id} ite {sort} {} {} {}",
            children[0], children[1], children[2]
        )?,

        Expr::ArrayStore { .. } => write!(
            writer,
            "{id} write {sort} {} {} {}",
            children[0], children[1], children[2]
        )?,

        Expr::ArrayConstant { .. } => {
            // TODO: this situation could be handled by creating a new constant state.
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "cannot serialize ArrayConstant outside of a state's init position: \
                 core BTOR2 has no operator for constant arrays",
            ));
        }
    }
    writeln!(writer, "{tail}")
}

fn write_bv_literal<W: Write>(
    writer: &mut W,
    ctx: &Context,
    id: u64,
    sort: u64,
    value: BVLitValue,
) -> std::io::Result<()> {
    // The body is written without a trailing newline; `write_node` appends the
    // name suffix (if any) and the newline.
    let v = value.get(ctx);
    if v.is_zero() {
        write!(writer, "{id} zero {sort}")
    } else if v.is_one() {
        write!(writer, "{id} one {sort}")
    } else if v.is_all_ones() {
        write!(writer, "{id} ones {sort}")
    } else {
        write!(writer, "{id} const {sort} {}", v.to_bit_str())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_is_is_autogen_name() {
        assert!(!is_autogen_name(""));
        assert!(is_autogen_name("_input"));
        assert!(!is_autogen_name("_input_"));
        assert!(is_autogen_name("_input_0"));
        assert!(is_autogen_name("_input_1"));
        assert!(is_autogen_name("_input_10"));
        assert!(is_autogen_name("_input_999999"));
        assert!(!is_autogen_name("_input_999_999"));
        assert!(!is_autogen_name("_input_999999_"));
    }
}