Computation Graph Must Call solve_witness() Before Serialization

[hczphn]

Description

When saving the computation graph, the current workflow requires calling ctx.solve_witness() before serialization.
Otherwise, reloading the graph later fails with the following error:

assertion left == right failed: Computation graph is already done, please compile or load it only once.
left: ComputationGraphDone
right: ComputationGraphNotDone

Reproduction

use expander_compiler::frontend::*;
use expander_compiler::zkcuda::proving_system::expander::config::{
    ZKCudaBN254KZG, ZKCudaBN254KZGBatchPCS, ZKCudaBN254MIMCKZGBatchPCS,
};
use expander_compiler::zkcuda::proving_system::expander_pcs_defered::BN254ConfigSha2UniKZG;
use expander_compiler::zkcuda::proving_system::{
    Expander, ExpanderNoOverSubscribe, ParallelizedExpander, ProvingSystem,
};
use expander_compiler::zkcuda::shape::Reshape;
use expander_compiler::zkcuda::{context::*, kernel::*};

use gkr::BN254ConfigSha2Hyrax;
use serdes::ExpSerde;
#[test]
fn zkcuda_assertion() {
    let kernel_tmp: KernelPrimitive<M31Config> = compile_assertion().unwrap();

    let mut ctx: Context<M31Config> = Context::default();
    let a = ctx.copy_to_device(&M31::from(10u32)).reshape(&[1]);
    let b = ctx.copy_to_device(&M31::from(10u32)).reshape(&[1]);
    call_kernel!(ctx, kernel_tmp, 1, a, b).unwrap();

    type P = Expander<M31Config>;
    let computation_graph1 = ctx.compile_computation_graph().unwrap();
    let file = std::fs::File::create("graph.txt").unwrap();
    let writer = std::io::BufWriter::new(file);
    computation_graph1.serialize_into(writer);

    let graph_data = std::fs::read("graph.txt").unwrap();
    let reader = std::io::BufReader::new(&graph_data[..]);
    let computation_graph: ComputationGraph<M31Config> =
        ComputationGraph::deserialize_from(reader).unwrap();

    ctx.load_computation_graph(computation_graph.clone()).unwrap();
    ctx.solve_witness().unwrap(); // ❌ Causes panic

    let (prover_setup, verifier_setup) = P::setup(&computation_graph);
    let proof = P::prove(
        &prover_setup,
        &computation_graph,
        ctx.export_device_memories(),
    );
    assert!(P::verify(&verifier_setup, &computation_graph, &proof));
}

Observed Behavior

If ctx.solve_witness() is not called before saving, deserialization fails because the graph state doesn’t match expectations.

If it is called before saving, the graph becomes tied to that specific witness, and there’s no API to update it later.

Expected Behavior

Users should be able to:

Serialize a computation graph before solving its witness.

Reload the graph later and supply a new witness (without recompiling).

Proposed Enhancement

Add an interface that allows providing new witness data after loading a graph, or directly when proving, for example:

let proof = P::prove(
    &prover_setup,
    &computation_graph,
    ctx.export_device_memories_with_new_witness(new_witness),
);

Summary

Currently, a computation graph can only be used with one witness due to this limitation.
Enhancing the API to support providing or updating witness data independently would allow reusing the same computation graph across multiple proof generations.