Restoring checkpoint into unchanged circuit silently loses all state inside recursive() scopes #6765

use dbsp::circuit::{
    CircuitConfig, CircuitStorageConfig, Mode, StorageCacheConfig, StorageConfig, StorageOptions,
};
use dbsp::typed_batch::SpineSnapshot;
use dbsp::utils::Tup2;
use dbsp::{IndexedZSetReader, OrdZSet, OutputHandle, Runtime, Stream, ZSetHandle};
use uuid::Uuid;

type Out = OutputHandle<SpineSnapshot<OrdZSet<i32>>>;

fn config(dir: &str, init: Option<Uuid>) -> CircuitConfig {
    let storage = CircuitStorageConfig::for_config(
        StorageConfig {
            path: dir.to_string(),
            cache: StorageCacheConfig::default(),
        },
        StorageOptions::default(),
    )
    .unwrap()
    .with_init_checkpoint(init);
    CircuitConfig::with_workers(1)
        .with_mode(Mode::Persistent)
        .with_storage(Some(storage))
}

fn build(cfg: CircuitConfig, extra_output: bool) -> (dbsp::DBSPHandle, ZSetHandle<i32>, Out) {
    let (handle, (input, output)) = Runtime::init_circuit(cfg, move |root| {
        let (numbers, input) = root.add_input_zset::<i32>();
        numbers.set_persistent_id(Some("numbers"));
        numbers.integrate_trace();
        let set = root
            .recursive(|child, rec: Stream<_, OrdZSet<i32>>| Ok(numbers.delta0(child).plus(&rec)))
            .unwrap();
        let output = set.accumulate_output_persistent(Some("set"));
        if extra_output {
            // A NEW stateful node; its missing checkpoint file forces need_backfill
            let _extra = set.accumulate_output_persistent(Some("set2"));
        }
        Ok((input, output))
    })
    .unwrap();
    (handle, input, output)
}

fn drain(output: &Out) -> Vec<(i32, i64)> {
    let Some(snapshot) = output.take_from_worker(0) else {
        return Vec::new();
    };
    let mut v: Vec<(i32, i64)> = snapshot.consolidate().iter().map(|(n, (), w)| (n, w)).collect();
    v.sort();
    v
}

fn main() {
    let modified = std::env::args().any(|a| a == "--modified");
    let oracle = std::env::args().any(|a| a == "--oracle");
    let dir = tempfile::tempdir().unwrap().keep().to_string_lossy().into_owned();

    // Feed {1}.
    let (mut handle, input, output) = build(config(&dir, None), false);
    input.append(&mut vec![Tup2(1, 1i64)]);
    handle.transaction().unwrap();
    let delta = drain(&output);
    println!("phase 1 delta from {{1}}: {delta:?}");
    assert_eq!(delta, vec![(1, 1)]);

    // Oracle keeps running; other arms checkpoint, kill, restore.
    let (mut handle, input, output) = if oracle {
        (handle, input, output)
    } else {
        let meta = handle.checkpoint().run().unwrap();
        handle.kill().unwrap();
        let (mut handle, input, output) = build(config(&dir, Some(meta.uuid)), modified);
        let mut boot = 0u32;
        while handle.bootstrap_in_progress() {
            handle.transaction().unwrap();
            boot += 1;
            assert!(boot <= 100, "bootstrap did not converge");
        }
        println!("bootstrap transactions: {boot}");
        (handle, input, output)
    };

    // Feed {1 (already present), 2 (new)}.
    input.append(&mut vec![Tup2(1, 1i64), Tup2(2, 1)]);
    handle.transaction().unwrap();
    let delta = drain(&output);
    println!("phase 2 delta from {{1, 2}}: {delta:?}");
    if oracle {
        return;
    }
    // [(2, 1)] is what the --oracle arm prints for this feed.
    if delta == vec![(2, 1)] {
        println!("PASS: recursive-scope state round-tripped");
    } else {
        println!("FAIL: expected [(2, 1)]; got {delta:?}");
        std::process::exit(1);
    }
}