// Copyright The Hush Developers  2019-2022
// Released under the GPLv3

//! Witness-history tests for the Tier-6 g2 changes (audit C232/C233).
//!
//! `ChainModel` is the OLD (pre-g2) witness algorithm, written independently of the wallet:
//! per block, every note's newest witness is duplicated and its list trimmed to 100 entries;
//! then each leaf of the block is appended to the newest witness of every unspent note, to the
//! witnesses of our notes created earlier in the same block, and to the tree. A k-block
//! invalidate pops k entries from every list. The wallet must equal the model entry for entry,
//! with one declared exception: a note whose confirmed spender is below a full
//! (MAX_REORG + 1)-block window may lose its history, and then must have lost it.
//!
//! `check_anchor_alignment` (T0) is the spend-validity oracle: the k-th newest witness of
//! every live note must have the root of the k-th newest retained block, which is exactly the
//! anchor a spend at anchor offset k proves against.

use std::cmp;
use std::convert::TryInto;

use ff::{Field, PrimeField, PrimeFieldRepr};
use pairing::bls12_381::{Bls12, Fr};
use protobuf::Message;
use rand::{rngs::OsRng, RngCore};

use zcash_client_backend::proto::compact_formats::{CompactBlock, CompactOutput, CompactSpend, CompactTx};
use zcash_primitives::{
    block::BlockHash,
    jubjub::fs::Fs,
    merkle_tree::{CommitmentTree, IncrementalWitness},
    note_encryption::{Memo, SaplingNoteEncryption},
    primitives::Note,
    sapling::Node,
    transaction::TxId,
    zip32::{ExtendedFullViewingKey, ExtendedSpendingKey},
    JUBJUB,
};

use super::data::SaplingNoteData;
use super::{LightWallet, MAX_REORG};
use crate::lightclient::LightClientConfig;

fn cfg() -> LightClientConfig {
    LightClientConfig {
        server: "0.0.0.0:0".parse().unwrap(),
        chain_name: "main".to_string(), // "test" panics in base58_pubkey_address
        sapling_activation_height: 0,
        consensus_branch_id: "000000".to_string(),
        anchor_offset: 0,
        no_cert_verification: false,
        data_dir: None, // a LightWallet never touches the disk itself
    }
}

fn wbytes(w: &IncrementalWitness<Node>) -> Vec<u8> {
    let mut v = vec![];
    w.write(&mut v).unwrap();
    v
}

fn tbytes(t: &CommitmentTree<Node>) -> Vec<u8> {
    let mut v = vec![];
    t.write(&mut v).unwrap();
    v
}

fn random_hash() -> [u8; 32] {
    let mut h = [0u8; 32];
    OsRng.fill_bytes(&mut h);
    h
}

fn foreign() -> ExtendedFullViewingKey {
    ExtendedFullViewingKey::from(&ExtendedSpendingKey::master(&[7u8; 32]))
}

/// An output paying `value` to the default address of `fvk` (the body of tests.rs add_tx_paying_to).
fn out_to(fvk: &ExtendedFullViewingKey, value: u64) -> CompactOutput {
    let mut rng = OsRng;
    let to = fvk.default_address().unwrap().1;
    let note = Note {
        g_d: to.diversifier.g_d::<Bls12>(&JUBJUB).unwrap(),
        pk_d: to.pk_d.clone(),
        value,
        r: Fs::random(&mut rng),
    };
    let encryptor = SaplingNoteEncryption::new(fvk.fvk.ovk, note.clone(), to.clone(), Memo::default(), &mut rng);
    let mut cmu = vec![];
    note.cm(&JUBJUB).into_repr().write_le(&mut cmu).unwrap();
    let mut epk = vec![];
    encryptor.epk().write(&mut epk).unwrap();
    let enc_ciphertext = encryptor.encrypt_note_plaintext();

    let mut cout = CompactOutput::new();
    cout.set_cmu(cmu);
    cout.set_epk(epk);
    cout.set_ciphertext(enc_ciphertext[..52].to_vec());
    cout
}

fn random_cmu_bytes() -> Vec<u8> {
    let mut v = vec![];
    Fr::random(&mut OsRng).into_repr().write_le(&mut v).unwrap();
    v
}

/// A valid cmu with a 31-byte epk: the leaf is committed, decryption is skipped.
fn bad_epk() -> CompactOutput {
    let mut o = CompactOutput::new();
    o.set_cmu(random_cmu_bytes());
    o.set_epk(vec![1u8; 31]);
    o.set_ciphertext(vec![0u8; 52]);
    o
}

/// A 31-byte cmu: nothing is appended anywhere.
fn bad_cmu() -> CompactOutput {
    let mut o = CompactOutput::new();
    o.set_cmu(vec![1u8; 31]);
    o.set_epk(vec![0u8; 32]);
    o.set_ciphertext(vec![0u8; 52]);
    o
}

fn tx_with(outputs: Vec<CompactOutput>, spend_nfs: Vec<Vec<u8>>) -> (CompactTx, TxId) {
    let hash = random_hash();
    let mut ctx = CompactTx::new();
    ctx.set_hash(hash.to_vec());
    for nf in spend_nfs {
        let mut s = CompactSpend::new();
        s.set_nf(nf);
        ctx.spends.push(s);
    }
    for o in outputs {
        ctx.outputs.push(o);
    }
    (ctx, TxId(hash))
}

/// A deterministic schedule (the note contents themselves use OsRng).
struct Lcg(u64);

impl Lcg {
    fn below(&mut self, n: u64) -> u64 {
        self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
        (self.0 >> 33) % n
    }
}

/// A tx for `Chain::mine`: its compact form, which of its outputs pays the wallet, and the
/// creating txids of the model notes it spends.
struct GenTx {
    ctx: CompactTx,
    txid: TxId,
    ours: Option<usize>,
    spends: Vec<TxId>,
}

struct ModelNote {
    /// Creating txid (every test tx pays the wallet at most one note).
    txid: TxId,
    recv: i32,
    /// Block of the spending tx.
    spent_at: Option<i32>,
    hist: Vec<IncrementalWitness<Node>>,
}

/// The OLD witness algorithm (see the module comment).
struct ChainModel {
    tree: CommitmentTree<Node>,
    /// Retained window: (height, tree at the end of that block), at most MAX_REORG + 1 like the wallet.
    win: Vec<(i32, CommitmentTree<Node>)>,
    notes: Vec<ModelNote>,
    /// A spent note whose spender block is below this may (and must) have lost its history:
    /// the max first height of a full window seen by a scan (before the block) or by a read.
    clear_bound: i32,
}

impl ChainModel {
    fn new() -> Self {
        ChainModel { tree: CommitmentTree::new(), win: vec![], notes: vec![], clear_bound: i32::MIN }
    }

    fn note(&self, txid: &TxId) -> &ModelNote {
        self.notes.iter().find(|n| n.txid == *txid).expect("model note")
    }

    fn note_full_window(&mut self) {
        if self.win.len() > MAX_REORG {
            self.clear_bound = cmp::max(self.clear_bound, self.win[0].0);
        }
    }

    fn must_be_cleared(&self, n: &ModelNote) -> bool {
        n.spent_at.map_or(false, |s| s < self.clear_bound)
    }

    /// nodes: the block's leaves in order (outputs whose cmu parses); ours: (index into nodes,
    /// creating txid) of the outputs paying the wallet; spends: creating txids of spent notes.
    fn apply_block(&mut self, height: i32, nodes: &[Node], ours: &[(usize, TxId)], spends: &[TxId]) {
        for n in self.notes.iter_mut() {
            if let Some(w) = n.hist.last() {
                let c = w.clone();
                n.hist.push(c);
            }
            n.hist = n.hist.split_off(n.hist.len().saturating_sub(100));
        }
        let active: Vec<usize> = (0..self.notes.len()).filter(|i| self.notes[*i].spent_at.is_none()).collect();
        let mut fresh: Vec<ModelNote> = vec![];
        for (k, node) in nodes.iter().enumerate() {
            for i in &active {
                if let Some(w) = self.notes[*i].hist.last_mut() {
                    w.append(*node).unwrap();
                }
            }
            for f in fresh.iter_mut() {
                f.hist[0].append(*node).unwrap();
            }
            self.tree.append(*node).unwrap();
            if let Some((_, txid)) = ours.iter().find(|(idx, _)| *idx == k) {
                fresh.push(ModelNote {
                    txid: *txid,
                    recv: height,
                    spent_at: None,
                    hist: vec![IncrementalWitness::from_tree(&self.tree)],
                });
            }
        }
        for t in spends {
            for n in self.notes.iter_mut() {
                if n.txid == *t {
                    n.spent_at = Some(height);
                }
            }
        }
        self.notes.extend(fresh);
        self.win.push((height, self.tree.clone()));
        if self.win.len() > MAX_REORG + 1 {
            let d = self.win.len() - (MAX_REORG + 1);
            self.win.drain(..d);
        }
    }

    /// OLD invalidate_block: pop the retained blocks at or above `at`, drop the notes created
    /// there, un-spend the notes spent there, and pop the same count from every list.
    fn invalidate(&mut self, at: i32) -> u64 {
        let mut k = 0usize;
        while self.win.last().map_or(false, |(h, _)| *h >= at) {
            self.win.pop();
            k += 1;
        }
        self.tree = match self.win.last() {
            Some((_, t)) => t.clone(),
            None => CommitmentTree::new(),
        };
        self.notes.retain(|n| n.recv < at);
        for n in self.notes.iter_mut() {
            if n.spent_at.map_or(false, |s| s >= at) {
                n.spent_at = None;
            }
            let keep = n.hist.len().saturating_sub(k);
            n.hist.truncate(keep);
        }
        k as u64
    }
}

struct Chain {
    wallet: LightWallet,
    model: ChainModel,
    prev: BlockHash,
    next: i32,
}

impl Chain {
    fn new() -> Self {
        Chain {
            wallet: LightWallet::new(None, &cfg(), 0, 0).unwrap(),
            model: ChainModel::new(),
            prev: BlockHash([0u8; 32]),
            next: 0,
        }
    }

    fn ours(&self) -> ExtendedFullViewingKey {
        self.wallet.zkeys.read().unwrap()[0].extfvk.clone()
    }

    /// The stored nullifier of the wallet note created by `txid` (its true tree position).
    fn nf(&self, txid: &TxId) -> Vec<u8> {
        self.wallet.txs.read().unwrap()[txid].notes[0].nullifier.to_vec()
    }

    fn pay_us(&self, value: u64) -> GenTx {
        let (ctx, txid) = tx_with(vec![out_to(&self.ours(), value)], vec![]);
        GenTx { ctx, txid, ours: Some(0), spends: vec![] }
    }

    fn pay_foreign(&self, n: usize) -> GenTx {
        let f = foreign();
        let (ctx, txid) = tx_with((0..n).map(|i| out_to(&f, 5_000 + i as u64)).collect(), vec![]);
        GenTx { ctx, txid, ours: None, spends: vec![] }
    }

    /// Spend the note created by `note`: [foreign payment, change to us].
    fn spend(&self, note: TxId, change: u64) -> GenTx {
        let (ctx, txid) = tx_with(vec![out_to(&foreign(), 7_000), out_to(&self.ours(), change)], vec![self.nf(&note)]);
        GenTx { ctx, txid, ours: Some(1), spends: vec![note] }
    }

    fn mine(&mut self, gtxs: Vec<GenTx>) -> CompactBlock {
        let height = self.next;
        let mut cb = CompactBlock::new();
        cb.set_height(height as u64);
        cb.hash = random_hash().to_vec();
        cb.prevHash = self.prev.0.to_vec();
        let mut nodes: Vec<Node> = vec![];
        let mut ours: Vec<(usize, TxId)> = vec![];
        let mut spends: Vec<TxId> = vec![];
        for (i, g) in gtxs.into_iter().enumerate() {
            for (j, o) in g.ctx.outputs.iter().enumerate() {
                if let Ok(cmu) = o.cmu() {
                    if g.ours == Some(j) {
                        ours.push((nodes.len(), g.txid));
                    }
                    nodes.push(Node::new(cmu.into_repr()));
                }
            }
            spends.extend(g.spends.iter().cloned());
            let mut ctx = g.ctx;
            ctx.set_index(i as u64);
            cb.vtx.push(ctx);
        }
        self.model.note_full_window();
        let r = self.wallet.scan_block(&cb.write_to_bytes().unwrap());
        assert!(r.is_ok(), "scan_block({}) failed: {:?}", height, r);
        self.model.apply_block(height, &nodes, &ours, &spends);
        self.prev = BlockHash(cb.hash[..].try_into().unwrap());
        self.next += 1;
        cb
    }

    fn rewind(&mut self) {
        self.next = self.wallet.last_scanned_height() + 1;
        self.prev = self.wallet.blocks.read().unwrap().last().map(|b| b.hash).unwrap_or(BlockHash([0u8; 32]));
    }

    fn invalidate(&mut self, at: i32) -> u64 {
        let n = self.wallet.invalidate_block(at);
        let m = self.model.invalidate(at);
        assert_eq!(n, m, "invalidate_block({}) count vs the model", at);
        self.rewind();
        n
    }

    /// Replace the wallet by `read(write(wallet))`. Returns (first write, write of the re-read wallet).
    fn reload(&mut self) -> (Vec<u8>, Vec<u8>) {
        let mut b1 = vec![];
        self.wallet.write(&mut b1).unwrap();
        self.wallet = LightWallet::read(&b1[..], &cfg()).unwrap();
        self.model.note_full_window();
        let mut b2 = vec![];
        self.wallet.write(&mut b2).unwrap();
        (b1, b2)
    }

    fn list_len(&self, txid: &TxId) -> usize {
        self.wallet.txs.read().unwrap()[txid].notes[0].witnesses.len()
    }

    /// Compare the wallet with the model: the block window and trees, every note's spent state,
    /// and its witness list (`full`: every entry's bytes; otherwise the count and the newest entry).
    /// Returns how many notes have the history the model requires to be dropped.
    fn compare(&self, full: bool, ctx: &str) -> usize {
        {
            let blocks = self.wallet.blocks.read().unwrap();
            let hs: Vec<i32> = blocks.iter().map(|b| b.height).collect();
            let mh: Vec<i32> = self.model.win.iter().map(|(h, _)| *h).collect();
            assert_eq!(hs, mh, "{}: retained block window", ctx);
            if let Some(b) = blocks.last() {
                assert!(tbytes(&b.tree) == tbytes(&self.model.tree), "{}: tip tree differs from the model", ctx);
            }
            if full {
                for (b, (h, t)) in blocks.iter().zip(self.model.win.iter()) {
                    assert!(tbytes(&b.tree) == tbytes(t), "{}: tree of block {} differs from the model", ctx, h);
                }
            }
        }
        let txs = self.wallet.txs.read().unwrap();
        let total: usize = txs.values().map(|t| t.notes.len()).sum();
        assert_eq!(total, self.model.notes.len(), "{}: number of wallet notes", ctx);
        let mut cleared = 0;
        for m in &self.model.notes {
            let wtx = txs.get(&m.txid).unwrap_or_else(|| panic!("{}: no wallet tx {}", ctx, m.txid));
            assert_eq!(wtx.notes.len(), 1, "{}: notes of tx {}", ctx, m.txid);
            assert_eq!(wtx.block, m.recv, "{}: block of tx {}", ctx, m.txid);
            let nd = &wtx.notes[0];
            match (nd.spent, m.spent_at) {
                (None, None) => {}
                (Some(s), Some(h)) => {
                    assert_eq!(txs.get(&s).map(|t| t.block), Some(h), "{}: spender block of note {}", ctx, m.txid)
                }
                (a, b) => panic!("{}: note {} spent {:?}, model spent at {:?}", ctx, m.txid, a, b),
            }
            assert!(nd.unconfirmed_spent.is_none(), "{}: note {} pending", ctx, m.txid);
            let l = nd.witnesses.len();
            assert!(l <= MAX_REORG, "{}: note {} has {} witnesses", ctx, m.txid, l);
            if self.model.must_be_cleared(m) {
                assert!(l == 0, "{}: note {} (spent at {:?}, clear bound {}) still has {} witnesses",
                        ctx, m.txid, m.spent_at, self.model.clear_bound, l);
                cleared += 1;
                continue;
            }
            assert_eq!(l, m.hist.len(), "{}: witness count of note {} (received {}, spent at {:?})",
                       ctx, m.txid, m.recv, m.spent_at);
            let from = if full { 0 } else { l.saturating_sub(1) };
            for i in from..l {
                assert!(wbytes(&nd.witnesses[i]) == wbytes(&m.hist[i]),
                        "{}: witness {} of {} of note {} (received {}, spent at {:?}) differs from the model",
                        ctx, i, l, m.txid, m.recv, m.spent_at);
            }
        }
        cleared
    }
}

/// T0: for every unspent note, pending ones included (L entries, B retained blocks) assert L <= 100,
/// L <= B and, for k in {0, 1, L-1}, witnesses[L-1-k].root() == blocks[B-1-k].tree.root(): the
/// anchor at offset k.
fn check_anchor_alignment(w: &LightWallet, ctx: &str) -> usize {
    let blocks = w.blocks.read().unwrap();
    let b = blocks.len();
    let txs = w.txs.read().unwrap();
    let mut checked = 0;
    for tx in txs.values() {
        for nd in tx.notes.iter() {
            if nd.spent.is_some() {
                continue;
            }
            let l = nd.witnesses.len();
            assert!(l <= MAX_REORG && l <= b, "{}: live note of tx {} has {} witnesses, {} blocks", ctx, tx.txid, l, b);
            let mut ks = vec![0usize, 1, l.saturating_sub(1)];
            ks.dedup();
            for k in ks {
                if k < l {
                    assert!(nd.witnesses[l - 1 - k].root() == blocks[b - 1 - k].tree.root(),
                            "{}: live note of tx {}: witness at offset {} is not the root of block {}",
                            ctx, tx.txid, k, blocks[b - 1 - k].height);
                    checked += 1;
                }
            }
        }
    }
    checked
}

/// Test 1's chain, heights 0..=259: block 0 pays note A; about 25% of blocks pay us 1-3 notes and
/// about 20% carry foreign outputs; at 40, 120 and 200 a tx spends the oldest live note (A first)
/// with [foreign, change to us]. With `check`, the wallet is compared with the model after every
/// block (all entries plus T0 and the roots at the checkpoints). Returns the (height, note) spends.
fn t1_chain(check: bool) -> (Chain, Vec<(i32, TxId)>) {
    let mut c = Chain::new();
    let mut rng = Lcg(0x5eed_0001);
    let mut spent: Vec<(i32, TxId)> = vec![];
    let checkpoints = [41, 100, 101, 141, 142, 143, 221, 222, 259];
    for h in 0..=259 {
        let mut txs = vec![];
        if h == 0 {
            txs.push(c.pay_us(100_000));
        }
        if h == 40 || h == 120 || h == 200 {
            let note = c.model.notes.iter()
                .find(|n| n.spent_at.is_none())
                .map(|n| n.txid)
                .expect("a live note to spend");
            txs.push(c.spend(note, 1_000));
            spent.push((h, note));
        }
        if h > 0 && rng.below(100) < 25 {
            for _ in 0..1 + rng.below(3) {
                let v = 1_000 + rng.below(1_000);
                txs.push(c.pay_us(v));
            }
        }
        if rng.below(100) < 20 {
            let n = 1 + rng.below(2) as usize;
            txs.push(c.pay_foreign(n));
        }
        c.mine(txs);
        if !check {
            continue;
        }
        let ctx = format!("t1 h={}", h);
        let full = checkpoints.contains(&h);
        c.compare(full, &ctx);
        // The clearing rule, independently of the model: a note spent at s is empty after
        // block h iff h >= 101 && s <= h - 102.
        for (s, note) in &spent {
            let expect_empty = h >= 101 && *s <= h - 102;
            assert_eq!(c.list_len(note) == 0, expect_empty, "{}: note spent at {} cleared?", ctx, s);
        }
        if full {
            check_anchor_alignment(&c.wallet, &ctx);
            let root = c.wallet.blocks.read().unwrap().last().unwrap().tree.root();
            assert!(root == c.model.tree.root(), "{}: tip root differs from the model", ctx);
        }
    }
    assert_eq!(spent[0].1, c.model.notes[0].txid, "the first spend must spend note A");
    (c, spent)
}

#[test]
fn witness_history_matches_old_algorithm() {
    let (c, spent) = t1_chain(true);
    assert_eq!(spent.iter().map(|(s, _)| *s).collect::<Vec<_>>(), vec![40, 120, 200]);
    assert_eq!(c.compare(true, "t1 end"), 2, "notes spent at 40 and 120 are cleared at 259, 200 is not");
    assert!(c.model.notes.len() > 40, "only {} notes", c.model.notes.len());
    let live_at_cap = c.model.notes.iter().filter(|n| n.spent_at.is_none() && n.hist.len() == MAX_REORG).count();
    assert!(live_at_cap > 10, "only {} live notes at the 100-entry cap", live_at_cap);
}

/// Test 2's chain: block 0 pays A, every 4th block pays us, some blocks carry foreign outputs,
/// block 10 spends A. Mined up to `tip`, compared with the model after every block; also pins
/// witness_clear_height() (None while <= 100 blocks are retained, then the window's first height)
/// and the block at which A is cleared (112).
fn t2_chain(tip: i32) -> (Chain, TxId) {
    let mut c = Chain::new();
    let mut a = None;
    for h in 0..=tip {
        let mut txs = vec![];
        if h == 0 {
            let g = c.pay_us(50_000);
            a = Some(g.txid);
            txs.push(g);
        }
        if h == 10 {
            txs.push(c.spend(a.unwrap(), 500));
        }
        if h % 4 == 1 {
            txs.push(c.pay_us(1_000 + h as u64));
        }
        if h % 7 == 3 {
            txs.push(c.pay_foreign(1));
        }
        c.mine(txs);
        let ctx = format!("t2 h={}", h);
        c.compare(h % 25 == 0 || h == tip, &ctx);
        let expect = if h >= MAX_REORG as i32 { Some(h - MAX_REORG as i32) } else { None };
        assert_eq!(c.wallet.witness_clear_height(), expect, "{}: witness_clear_height", ctx);
        if h >= 10 {
            assert_eq!(c.list_len(&a.unwrap()) == 0, h >= 112, "{}: note A cleared?", ctx);
        }
    }
    (c, a.unwrap())
}

#[test]
fn clear_then_reorg_matches_old_algorithm() {
    // (a) Tip 60: the window (61 blocks) is not full, nothing is cleared.
    let (mut c, a) = t2_chain(60);
    assert_eq!(c.compare(true, "2a tip 60"), 0);
    assert_eq!(c.list_len(&a), 61);
    assert_eq!(c.wallet.clear_old_witnesses(), 0);
    assert_eq!(c.invalidate(10), 51);
    assert!(c.wallet.txs.read().unwrap()[&a].notes[0].spent.is_none(), "A is unspent again");
    assert_eq!(c.list_len(&a), 10);
    c.compare(true, "2a after invalidate(10)");
    assert!(check_anchor_alignment(&c.wallet, "2a after invalidate(10)") > 0);
    // A new branch: the revived note advances like every other live note.
    for h in 0..6 {
        let txs = if h == 2 { vec![c.pay_us(4_242)] } else { vec![c.pay_foreign(1)] };
        c.mine(txs);
        c.compare(true, "2a new branch");
        check_anchor_alignment(&c.wallet, "2a new branch");
    }
    assert_eq!(c.list_len(&a), 16);

    // (b) Tip 115: A was cleared at 112. invalidate(10) pops the whole window; A is unspent and
    // empty, exactly like the old algorithm's trimmed list.
    let (mut c, a) = t2_chain(115);
    assert_eq!(c.compare(true, "2b tip 115"), 1);
    assert_eq!(c.invalidate(10), 101);
    assert!(c.wallet.txs.read().unwrap()[&a].notes[0].spent.is_none(), "A is unspent again");
    assert_eq!(c.list_len(&a), 0);
    assert!(c.model.note(&a).hist.is_empty(), "the old algorithm also holds an empty list");
    assert_eq!(c.compare(true, "2b after invalidate(10)"), 0);
    check_anchor_alignment(&c.wallet, "2b after invalidate(10)");

    // (c) Tip 115, a partial reorg: A stays spent and empty, live notes equal the model trimmed
    // by 36; then the rest of the window (65 blocks) is popped.
    let (mut c, a) = t2_chain(115);
    assert_eq!(c.invalidate(80), 36);
    assert!(c.wallet.txs.read().unwrap()[&a].notes[0].spent.is_some(), "A is still spent");
    assert_eq!(c.list_len(&a), 0);
    assert_eq!(c.compare(true, "2c after invalidate(80)"), 1);
    assert!(check_anchor_alignment(&c.wallet, "2c after invalidate(80)") > 0);
    assert_eq!(c.invalidate(10), 65);
    assert_eq!(c.list_len(&a), 0);
    assert!(c.model.note(&a).hist.is_empty());
    assert_eq!(c.compare(true, "2c after invalidate(10)"), 0);
}

#[test]
fn witness_serialization_compat() {
    // (a) write -> read -> write is byte-identical (clearing on read is idempotent).
    let (mut c, spent) = t1_chain(false);
    assert_eq!(c.compare(true, "3a built"), 2);
    let (b1, b2) = c.reload();
    assert!(b1 == b2, "write(read(write(w))) differs from write(w)");
    assert_eq!(c.wallet.clear_old_witnesses(), 0);
    assert_eq!(c.compare(true, "3a reloaded"), 2);
    check_anchor_alignment(&c.wallet, "3a reloaded");

    // (b) A legacy file: the cleared notes carry the old algorithm's full lists. Reading it
    // drops them again; everything else is unchanged, so the re-write equals the new file.
    let cleared: Vec<TxId> = spent.iter().map(|(_, t)| *t).filter(|t| c.list_len(t) == 0).collect();
    assert_eq!(cleared.len(), 2);
    for t in &cleared {
        let legacy = {
            let txs = c.wallet.txs.read().unwrap();
            let mut b = vec![];
            txs[t].notes[0].write(&mut b).unwrap();
            // version, account, extfvk (169), diversifier (11), value, rcm (32), then the witness Vector
            let at = 8 + 8 + 169 + 11 + 8 + 32;
            assert_eq!(b[at], 0, "an empty witness Vector (CompactSize 0) at offset {}", at);
            let hist = &c.model.note(t).hist;
            assert_eq!(hist.len(), MAX_REORG);
            let mut v = b[..at].to_vec();
            v.push(hist.len() as u8); // CompactSize below 253 is one byte
            for w in hist.iter() {
                v.extend(wbytes(w));
            }
            v.extend_from_slice(&b[at + 1..]);
            v
        };
        let nd = SaplingNoteData::read(&legacy[..]).unwrap();
        assert_eq!(nd.witnesses.len(), MAX_REORG);
        let mut again = vec![];
        nd.write(&mut again).unwrap();
        assert!(again == legacy, "the legacy note record does not round-trip");
        c.wallet.txs.write().unwrap().get_mut(t).unwrap().notes[0] = nd;
        assert_eq!(c.list_len(t), MAX_REORG);
    }
    let (legacy_file, rewritten) = c.reload();
    for t in &cleared {
        assert_eq!(c.list_len(t), 0, "reading a legacy file clears the long-spent note again");
    }
    assert!(rewritten == b1, "the legacy file re-writes to exactly the new file");
    assert!(rewritten.len() < legacy_file.len(), "{} vs {}", rewritten.len(), legacy_file.len());
    assert_eq!(c.compare(true, "3b reloaded"), 2);

    // (c) A 50-block chain (window not full): the spent note keeps its list through a reload.
    let (mut c50, a) = {
        let mut c = Chain::new();
        let mut a = None;
        for h in 0..50 {
            let mut txs = vec![];
            if h == 0 {
                let g = c.pay_us(30_000);
                a = Some(g.txid);
                txs.push(g);
            }
            if h == 10 {
                txs.push(c.spend(a.unwrap(), 300));
            }
            if h % 3 == 2 {
                txs.push(c.pay_us(2_000 + h as u64));
            }
            c.mine(txs);
        }
        (c, a.unwrap())
    };
    assert_eq!(c50.list_len(&a), 50);
    let (x1, x2) = c50.reload();
    assert!(x1 == x2);
    assert_eq!(c50.list_len(&a), 50);
    assert_eq!(c50.compare(true, "3c reloaded"), 0);

    // (d) A pending note and a note spent by a tx the wallet does not know are never cleared.
    let (x, y) = {
        let live: Vec<TxId> = c.model.notes.iter().filter(|n| n.spent_at.is_none()).map(|n| n.txid).take(2).collect();
        (live[0], live[1])
    };
    let before = |c: &Chain, t: &TxId| -> Vec<Vec<u8>> {
        c.wallet.txs.read().unwrap()[t].notes[0].witnesses.iter().map(|w| wbytes(w)).collect()
    };
    let (bx, by) = (before(&c, &x), before(&c, &y));
    assert!(!bx.is_empty() && !by.is_empty());
    {
        let mut txs = c.wallet.txs.write().unwrap();
        txs.get_mut(&x).unwrap().notes[0].unconfirmed_spent = Some(TxId(random_hash()));
        txs.get_mut(&y).unwrap().notes[0].spent = Some(TxId(random_hash()));
    }
    assert!(c.wallet.witness_clear_height().is_some(), "the window is full");
    assert_eq!(c.wallet.clear_old_witnesses(), 0);
    assert!(before(&c, &x) == bx && before(&c, &y) == by, "pending / unknown-spender notes were touched");
}

fn random_node() -> Node {
    Node::new(Fr::random(&mut OsRng).into_repr())
}

#[test]
fn append_nodes_parallel_matches_serial() {
    for &k in &[0usize, 1, 7, 8, 16, 17, 37, 64] {
        // K witnesses at different positions of one tree, all following it to its tip.
        let mut tree = CommitmentTree::<Node>::new();
        let mut base: Vec<IncrementalWitness<Node>> = vec![];
        for i in 0..k {
            for _ in 0..1 + i % 3 {
                let n = random_node();
                tree.append(n).unwrap();
                for w in base.iter_mut() {
                    w.append(n).unwrap();
                }
            }
            base.push(IncrementalWitness::from_tree(&tree));
        }
        for &n_nodes in &[0usize, 1, 5] {
            let ctx = format!("K={} nodes={}", k, n_nodes);
            let nodes: Vec<Node> = (0..n_nodes).map(|_| random_node()).collect();
            let mut tip = tree.clone();
            for n in &nodes {
                tip.append(*n).unwrap();
            }
            // (i) a literal double loop
            let mut serial = base.clone();
            for w in serial.iter_mut() {
                for n in &nodes {
                    w.append(*n).unwrap();
                }
            }
            // (ii) the helper, forced inline
            let mut inline = base.clone();
            {
                let mut refs: Vec<&mut IncrementalWitness<Node>> = inline.iter_mut().collect();
                LightWallet::append_nodes_to_witnesses(&mut refs[..], &nodes, usize::MAX);
            }
            // (iii) the helper with no size threshold (threaded for K >= 16 on >= 2 cores)
            let mut par = base.clone();
            {
                let mut refs: Vec<&mut IncrementalWitness<Node>> = par.iter_mut().collect();
                LightWallet::append_nodes_to_witnesses(&mut refs[..], &nodes, 0);
            }
            let root = tip.root();
            for i in 0..k {
                let s = wbytes(&serial[i]);
                assert!(s == wbytes(&inline[i]), "{}: witness {} differs (inline)", ctx, i);
                assert!(s == wbytes(&par[i]), "{}: witness {} differs (parallel)", ctx, i);
                assert!(par[i].root() == root, "{}: witness {} is not at the tree's root", ctx, i);
            }
        }
    }
}

/// K notes, one per block; then one block whose first tx pays us and carries an output with an
/// unparseable epk (its leaf is committed) and one with an unparseable cmu (nothing is appended),
/// and whose second tx carries foreign outputs: 8 leaves. With K = 17 that is 17 x 8 = 136
/// existing-witness appends, which takes the deferred, threaded path; K = 3 takes the inline one.
fn deferred_case(k: usize) {
    let mut c = Chain::new();
    for _ in 0..k {
        let g = c.pay_us(10_000);
        c.mine(vec![g]);
    }
    let ctx = format!("deferred K={}", k);
    c.compare(true, &ctx);
    assert_eq!(c.model.notes.iter().filter(|n| n.spent_at.is_none()).count(), k);

    let f = foreign();
    let ours = c.ours();
    let (t1, id1) = tx_with(vec![out_to(&ours, 777), bad_epk(), bad_cmu(), out_to(&f, 1), out_to(&f, 2), out_to(&f, 3)], vec![]);
    let (t2, id2) = tx_with(vec![out_to(&f, 4), out_to(&f, 5), out_to(&f, 6)], vec![]);
    let size_before = c.model.tree.size();
    c.mine(vec![
        GenTx { ctx: t1, txid: id1, ours: Some(0), spends: vec![] },
        GenTx { ctx: t2, txid: id2, ours: None, spends: vec![] },
    ]);
    assert_eq!(c.model.tree.size() - size_before, 8, "{}: leaves of the bad-output block", ctx);
    c.compare(true, &ctx);
    assert!(check_anchor_alignment(&c.wallet, &ctx) >= k + 1, "{}: every live note is checked", ctx);
    let root = c.wallet.blocks.read().unwrap().last().unwrap().tree.root();
    assert!(root == c.model.tree.root(), "{}: tip root differs from the model", ctx);

    // One more block: every note, including the one created in the bad-output block, advances.
    let g = c.pay_foreign(2);
    c.mine(vec![g]);
    c.compare(true, &ctx);
    check_anchor_alignment(&c.wallet, &ctx);
}

#[test]
fn deferred_append_path_with_bad_outputs() {
    deferred_case(17);
    deferred_case(3);
}