pub mod valuation;
pub mod value;

use crate::analysis::cpa::lattice::JoinSemiLattice;
use crate::analysis::cpa::residue::EmptyResidue;
use crate::analysis::cpa::state::{AbstractState, MergeOutcome, Successor};
use crate::analysis::cpa::{ConfigurableProgramAnalysis, IntoState};
use crate::analysis::valuation::simple::valuation::ValuationSet;
use crate::analysis::varnode_map::VarNodeMap;
use crate::display::JingleDisplay;
use crate::modeling::machine::cpu::concrete::ConcretePcodeAddress;
use jingle_sleigh::{GeneralizedVarNode, PcodeOperation, SleighArchInfo, SpaceType, VarNode};
use std::borrow::Borrow;
use std::cmp::Ordering;
use std::fmt::{Display, Formatter, Result as FmtResult};
use std::hash::{Hash, Hasher};

use crate::analysis::valuation::simple::value::{Load, Value};
use std::rc::Rc;

/// How to merge conflicting valuations for a single varnode when joining states.
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
pub enum MergeBehavior {
    /// Combine differing valuations into a `Choice(...)` expression (higher precision).
    Choice,
    /// Converge differing valuations to `Top` (lower precision).
    Top,
}

/// State for the valuation CPA. Stores a `ValuationSet` which contains both
/// direct and indirect write maps.
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub struct ValuationState {
    valuation: ValuationSet,
    arch_info: SleighArchInfo,
    /// Merge behavior controlling how conflicting valuations are handled during `join`.
    merge_behavior: MergeBehavior,
}

impl AsRef<SleighArchInfo> for ValuationState {
    fn as_ref(&self) -> &SleighArchInfo {
        &self.arch_info
    }
}

impl Display for ValuationState {
    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
        use std::collections::hash_map::DefaultHasher;
        let mut hasher = DefaultHasher::new();
        self.hash(&mut hasher);
        let hash_value = hasher.finish();
        write!(f, "Hash({:016x})", hash_value)
    }
}

impl JingleDisplay for ValuationState {
    fn fmt_jingle(&self, f: &mut Formatter<'_>, info: &SleighArchInfo) -> std::fmt::Result {
        // Delegate display to the inner Valuation implementation to avoid duplication.
        self.valuation.fmt_jingle(f, info)
    }
}

impl ValuationState {
    /// Create a new state with the default merge behavior of `Or`.
    pub fn new(arch_info: SleighArchInfo) -> Self {
        Self {
            valuation: ValuationSet::new(),
            arch_info,
            merge_behavior: MergeBehavior::Choice,
        }
    }

    /// Create a new state specifying the desired merge behavior.
    pub fn new_with_behavior(arch_info: SleighArchInfo, merge_behavior: MergeBehavior) -> Self {
        Self {
            valuation: ValuationSet::new(),
            arch_info,
            merge_behavior,
        }
    }

    pub fn get_value(&self, varnode: &VarNode) -> Option<&Value> {
        self.valuation.direct_writes.get(varnode)
    }

    pub fn written_locations(&self) -> &VarNodeMap<Value> {
        &self.valuation.direct_writes
    }

    pub fn valuation(&self) -> &ValuationSet {
        &self.valuation
    }

    /// Transfer function: build symbolic valuations for pcode operations.
    /// This returns a new state (functional) instead of mutating in place.
    fn transfer_impl(&self, op: &PcodeOperation) -> Self {
        let mut new_state = self.clone();

        // Match on the operation. Handle stores (indirect) and direct-output ops.
        match op {
            // Store: record Load(ptr, size) -> value in indirect_writes
            PcodeOperation::Store { output, input } => {
                let ptr = &output.pointer_location();
                let val = Value::from_varnode_or_entry(self, input);
                let pv = Value::from_varnode_or_entry(self, ptr);
                let data_size = input.size();
                let loc = Value::Load(Load(
                    Rc::new(pv.simplify()),
                    data_size,
                    output.pointer_space_index() as u8,
                ));
                new_state.valuation.add(loc, val);
            }

            // Copy
            PcodeOperation::Copy { input, .. } => {
                let result = if input.is_const() {
                    Value::const_(input.offset() as i64, input.size())
                } else {
                    Value::from_varnode_or_entry(self, input)
                };
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, result);
                }
            }

            PcodeOperation::IntAdd { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, a + b);
                }
            }

            PcodeOperation::IntSub { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, a - b);
                }
            }

            PcodeOperation::IntXor { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, a ^ b);
                }
            }

            PcodeOperation::IntMult { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, a * b);
                }
            }

            PcodeOperation::IntOr { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, a | b);
                }
            }

            PcodeOperation::IntAnd { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, a & b);
                }
            }

            PcodeOperation::BoolNegate { input, .. } => {
                let a = Value::from_varnode_or_entry(self, input);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, Value::bool_negate(a));
                }
            }

            PcodeOperation::BoolAnd { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, Value::bool_and(a, b));
                }
            }

            PcodeOperation::BoolOr { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, Value::bool_or(a, b));
                }
            }

            PcodeOperation::BoolXor { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, Value::bool_xor(a, b));
                }
            }

            PcodeOperation::IntLeftShift { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    let s = std::cmp::max(a.size(), b.size());
                    let shift_expr =
                        Value::IntLeftShift(value::IntLeftShiftExpr(Rc::new(a), Rc::new(b), s));
                    new_state.valuation.add(output_vn, shift_expr);
                }
            }

            PcodeOperation::IntRightShift { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    let s = std::cmp::max(a.size(), b.size());
                    let shift_expr =
                        Value::IntRightShift(value::IntRightShiftExpr(Rc::new(a), Rc::new(b), s));
                    new_state.valuation.add(output_vn, shift_expr);
                }
            }

            PcodeOperation::IntSignedRightShift { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    let s = std::cmp::max(a.size(), b.size());
                    let shift_expr = Value::IntSignedRightShift(value::IntSignedRightShiftExpr(
                        Rc::new(a),
                        Rc::new(b),
                        s,
                    ));
                    new_state.valuation.add(output_vn, shift_expr);
                }
            }

            PcodeOperation::IntNegate { input, .. } => {
                let a = Value::const_(0, input.size());
                let b = Value::from_varnode_or_entry(self, input);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, a - b);
                }
            }

            PcodeOperation::Int2Comp { input, .. } => {
                let a = Value::from_varnode_or_entry(self, input);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, Value::int_2comp(a));
                }
            }

            PcodeOperation::Load { input, .. } => {
                let ptr = &input.pointer_location();
                let pv = Value::from_varnode_or_entry(self, ptr);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    let load_expr = Value::Load(Load(
                        Rc::new(pv.simplify()),
                        output_vn.size(),
                        input.pointer_space_index() as u8,
                    ));
                    if let Some(v) = self.valuation.indirect_writes.get(&load_expr) {
                        new_state.valuation.add(output_vn, v.clone());
                    } else {
                        new_state.valuation.add(output_vn, load_expr);
                    }
                }
            }

            PcodeOperation::IntZExt { input, .. } => {
                let v = Value::from_varnode_or_entry(self, input);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    let out_size = output_vn.size();
                    new_state
                        .valuation
                        .add(output_vn, Value::zero_extend(v, out_size));
                }
            }

            PcodeOperation::IntSExt { input, .. } => {
                let v = Value::from_varnode_or_entry(self, input);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    let out_size = output_vn.size();
                    new_state
                        .valuation
                        .add(output_vn, Value::sign_extend(v, out_size));
                }
            }

            PcodeOperation::SubPiece { input0, input1, .. } => {
                let v = Value::from_varnode_or_entry(self, input0);
                let byte_offset = input1.offset() as usize;
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    let out_size = output_vn.size();
                    new_state
                        .valuation
                        .add(output_vn, Value::extract(v, byte_offset, out_size));
                }
            }

            PcodeOperation::IntEqual { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, Value::int_equal(a, b));
                }
            }

            PcodeOperation::IntSignedLess { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, Value::int_sless(a, b));
                }
            }

            PcodeOperation::IntLess { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, Value::int_less(a, b));
                }
            }

            PcodeOperation::PopCount { input, .. } => {
                let a = Value::from_varnode_or_entry(self, input);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, Value::popcount(a));
                }
            }

            PcodeOperation::IntNotEqual { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state
                        .valuation
                        .add(output_vn, Value::int_not_equal(a, b));
                }
            }

            PcodeOperation::IntLessEqual { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state
                        .valuation
                        .add(output_vn, Value::int_less_equal(a, b));
                }
            }

            PcodeOperation::IntSignedLessEqual { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state
                        .valuation
                        .add(output_vn, Value::int_sless_equal(a, b));
                }
            }

            PcodeOperation::IntCarry { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, Value::int_carry(a, b));
                }
            }

            PcodeOperation::IntSignedCarry { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, Value::int_scarry(a, b));
                }
            }

            PcodeOperation::IntSignedBorrow { input0, input1, .. } => {
                let a = Value::from_varnode_or_entry(self, input0);
                let b = Value::from_varnode_or_entry(self, input1);
                if let Some(GeneralizedVarNode::Direct(output_vn)) = op.output() {
                    new_state.valuation.add(output_vn, Value::int_sborrow(a, b));
                }
            }

            // Other operations we don't model produce writes of Top.
            _ => {
                if let Some(GeneralizedVarNode::Direct(vn)) = op.output() {
                    // todo handle indirect
                    new_state.valuation.add(vn, Value::Top);
                }
            }
        }

        // Clear internal-space varnodes on control-flow to non-const destinations (same policy as direct_valuation.rs)
        match op {
            PcodeOperation::Branch { input }
            | PcodeOperation::CBranch { input0: input, .. }
            | PcodeOperation::Fallthrough { input }
                if !input.is_const() =>
            {
                // VarNodeMap doesn't provide `retain`; collect keys to remove and remove them.
                let mut to_remove: Vec<VarNode> = Vec::new();
                for (vn, _) in new_state.valuation.direct_writes.items() {
                    let keep = self
                        .arch_info
                        .get_space(vn.space_index())
                        .map(|s| s._type != SpaceType::IPTR_CONSTANT)
                        .unwrap_or(false);
                    if !keep {
                        to_remove.push(*vn);
                    }
                }
                for k in to_remove {
                    new_state.valuation.direct_writes.remove(k);
                }
            }
            PcodeOperation::BranchInd { input } | PcodeOperation::CallInd { input } => {
                // Clear IPTR_INTERNAL varnodes except the branch target, which must survive
                // so that strengthen_from_valuation can read it.
                let branch_target = input.pointer_location();
                let mut to_remove: Vec<VarNode> = Vec::new();
                for (vn, _) in new_state.valuation.direct_writes.items() {
                    if vn == branch_target {
                        continue;
                    }
                    let keep = self
                        .arch_info
                        .get_space(vn.space_index())
                        .map(|space| space._type != SpaceType::IPTR_INTERNAL)
                        .unwrap_or(true);
                    if !keep {
                        to_remove.push(*vn);
                    }
                }
                for k in to_remove {
                    new_state.valuation.direct_writes.remove(k);
                }
            }
            PcodeOperation::Call { .. } => {}
            PcodeOperation::CallOther {
                output: Some(a), ..
            } => {
                new_state.valuation.add(a, Value::Top);
            }
            _ => {}
        }

        new_state
    }
}

impl PartialOrd for ValuationState {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        // Make states comparable only when they have the same direct keys and identical valuations.
        if self.valuation.direct_writes.len() != other.valuation.direct_writes.len() {
            return None;
        }

        for (key, val) in self.valuation.direct_writes.items() {
            match other.valuation.direct_writes.get(key) {
                Some(other_val) => {
                    if val != other_val {
                        return None;
                    }
                }
                None => return None,
            }
        }

        // Also require indirect maps to be identical for comparability.
        if self.valuation.indirect_writes.len() != other.valuation.indirect_writes.len() {
            return None;
        }
        for (k, v) in &self.valuation.indirect_writes {
            match other.valuation.indirect_writes.get(k) {
                Some(ov) => {
                    if v != ov {
                        return None;
                    }
                }
                None => return None,
            }
        }

        Some(Ordering::Equal)
    }
}

impl JoinSemiLattice for ValuationState {
    fn join(&mut self, other: &Self) {
        // Merge direct writes
        for (key, other_val) in other.valuation.direct_writes.items() {
            match self.valuation.direct_writes.get_mut(key) {
                Some(my_val) => {
                    if my_val == &Value::Top || other_val == &Value::Top {
                        *my_val = Value::Top;
                    } else if my_val != other_val {
                        match self.merge_behavior {
                            MergeBehavior::Choice => {
                                let combined = Value::choice(my_val.clone(), other_val.clone());
                                *my_val = combined.simplify();
                            }
                            MergeBehavior::Top => {
                                *my_val = Value::Top;
                            }
                        }
                    }
                }
                None => {
                    match self.merge_behavior {
                        MergeBehavior::Choice => {
                            let entry = Value::from_varnode_or_entry(self, key);
                            let choice = Value::choice(entry, other_val.clone());
                            self.valuation.add(*key, choice);
                        }
                        MergeBehavior::Top => {
                            // If the other state has a direct write that we don't, we have to assume it could be anything.
                            self.valuation.add(*key, Value::Top);
                            continue;
                        }
                    }
                }
            }
        }

        // Merge indirect writes (pointer -> value)
        for (key, other_val) in &other.valuation.indirect_writes {
            match self.valuation.indirect_writes.get_mut(key) {
                Some(my_val) => {
                    if my_val == &Value::Top || other_val == &Value::Top {
                        *my_val = Value::Top;
                    } else if my_val != other_val {
                        match self.merge_behavior {
                            MergeBehavior::Choice => {
                                let combined = Value::choice(my_val.clone(), other_val.clone());
                                *my_val = combined.simplify();
                            }
                            MergeBehavior::Top => {
                                *my_val = Value::Top;
                            }
                        }
                    }
                }
                None => match self.merge_behavior {
                    MergeBehavior::Choice => {
                        let choice = Value::choice(key.clone(), other_val.clone());
                        self.valuation.add(key.clone(), choice);
                    }
                    MergeBehavior::Top => {
                        // One branch writes, the other doesn't — result is unknown.
                        self.valuation.add(key.clone(), Value::Top);
                    }
                },
            }
        }
    }
}

impl AbstractState for ValuationState {
    fn merge(&mut self, other: &Self) -> MergeOutcome {
        self.merge_join(other)
    }

    fn stop<'a, T: Iterator<Item = &'a Self>>(&'a self, states: T) -> bool {
        // Defer to the standard stop predicate helper if available.
        self.stop_sep(states)
    }

    fn transfer<'a, B: Borrow<PcodeOperation>>(&'a self, opcode: B) -> Successor<'a, Self> {
        let next_state = self.transfer_impl(opcode.borrow());
        std::iter::once(next_state).into()
    }
}

pub struct ValuationAnalysis {
    arch_info: SleighArchInfo,
    /// Default merge behavior for states produced by this analysis.
    merge_behavior: MergeBehavior,
}

impl ValuationAnalysis {
    /// Create with the default merge behavior (`Or`).
    pub fn new(arch_info: SleighArchInfo, merge_behavior: MergeBehavior) -> Self {
        Self {
            arch_info,
            merge_behavior,
        }
    }
}

impl ConfigurableProgramAnalysis for ValuationAnalysis {
    type State = ValuationState;
    type Reducer<'op> = EmptyResidue<Self::State>;
}

impl IntoState<ValuationAnalysis> for ConcretePcodeAddress {
    fn into_state(
        self,
        c: &ValuationAnalysis,
    ) -> <ValuationAnalysis as ConfigurableProgramAnalysis>::State {
        ValuationState {
            valuation: ValuationSet::new(),
            arch_info: c.arch_info.clone(),
            merge_behavior: c.merge_behavior,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use jingle_sleigh::{SleighEndianness, SpaceInfo, SpaceType};

    fn test_arch() -> SleighArchInfo {
        SleighArchInfo::new(
            "test:LE:64:default".to_string(),
            std::iter::empty(),
            vec![
                SpaceInfo {
                    name: "const".to_string(),
                    index: 0,
                    index_size_bytes: 8,
                    word_size_bytes: 1,
                    _type: SpaceType::IPTR_CONSTANT,
                    endianness: SleighEndianness::Little,
                },
                SpaceInfo {
                    name: "register".to_string(),
                    index: 1,
                    index_size_bytes: 8,
                    word_size_bytes: 1,
                    _type: SpaceType::IPTR_PROCESSOR,
                    endianness: SleighEndianness::Little,
                },
            ]
            .into_iter(),
            1,
            vec![],
        )
    }

    fn reg(offset: u64, size: usize) -> VarNode {
        VarNode::new(offset, size, 1)
    }

    #[test]
    fn transfer_bool_negate_folds_constant() {
        let state = ValuationState::new(test_arch());
        let out = reg(0x10, 1);
        let op = PcodeOperation::BoolNegate {
            output: out,
            input: VarNode::new_const(0, 1),
        };

        let next = state.transfer_impl(&op);

        assert_eq!(next.get_value(&out), Some(&Value::const_(1, 1)));
    }

    #[test]
    fn transfer_bool_and_uses_boolean_value_node() {
        let state = ValuationState::new(test_arch());
        let out = reg(0x10, 1);
        let left = reg(0x20, 8);
        let right = reg(0x30, 8);
        let op = PcodeOperation::BoolAnd {
            output: out,
            input0: left,
            input1: right,
        };

        let next = state.transfer_impl(&op);
        let value = next.get_value(&out).expect("expected output valuation");
        let expr = value.as_bool_and().expect("expected BoolAnd node");

        assert_eq!(expr.0.as_ref(), &Value::entry(left));
        assert_eq!(expr.1.as_ref(), &Value::entry(right));
    }

    #[test]
    fn transfer_bool_xor_folds_to_bool_negate_when_rhs_true() {
        let mut state = ValuationState::new(test_arch());
        let out = reg(0x10, 1);
        let flag = reg(0x20, 1);
        state.valuation.add(
            flag,
            Value::int_equal(Value::entry(reg(0x30, 8)), Value::entry(reg(0x40, 8))),
        );
        let op = PcodeOperation::BoolXor {
            output: out,
            input0: flag,
            input1: VarNode::new_const(1, 1),
        };

        let next = state.transfer_impl(&op);
        let value = next.get_value(&out).expect("expected output valuation");

        assert_eq!(
            *value,
            Value::int_not_equal(Value::entry(reg(0x30, 8)), Value::entry(reg(0x40, 8))).simplify()
        );
    }

    #[test]
    fn transfer_read_sub_register_after_wider_zext_write() {
        // EAX is written, then RAX = INT_ZEXT EAX (evicts EAX), then ZF = (EAX == 0).
        // ZF must reflect the computed EAX value, not Entry(EAX).
        let mut state = ValuationState::new(test_arch());
        let eax = reg(0x0, 4);
        let rax = reg(0x0, 8);
        let zf = reg(0x100, 1);
        let src = reg(0x200, 4);

        state
            .valuation
            .add(eax, Value::entry(src) + Value::const_(1, 4));

        // RAX = INT_ZEXT EAX — this evicts EAX from direct_writes
        state = state.transfer_impl(&PcodeOperation::IntZExt {
            input: eax,
            output: rax,
        });
        assert!(
            state.get_value(&eax).is_none(),
            "EAX should be evicted after RAX write"
        );

        // ZF = (EAX == 0)
        state = state.transfer_impl(&PcodeOperation::IntEqual {
            output: zf,
            input0: eax,
            input1: VarNode::new_const(0, 4),
        });

        let expected =
            Value::int_equal(Value::entry(src) + Value::const_(1, 4), Value::const_(0, 4))
                .simplify();
        assert_eq!(*state.get_value(&zf).expect("ZF should be set"), expected);
    }

    #[test]
    fn transfer_read_sub_register_nonzero_byte_offset() {
        // Write a 4-byte register at offset 0, then read the upper 2 bytes (offset 2).
        // The result should be extract(wider_val, 2, 2), not Entry(upper).
        let mut state = ValuationState::new(test_arch());
        let wide = reg(0x0, 4);
        let upper = reg(0x2, 2); // bytes [2..4] of wide
        let out = reg(0x100, 1);
        let src = reg(0x200, 4);

        state.valuation.add(wide, Value::entry(src));

        state = state.transfer_impl(&PcodeOperation::IntEqual {
            output: out,
            input0: upper,
            input1: VarNode::new_const(0, 2),
        });

        let expected =
            Value::int_equal(Value::extract(Value::entry(src), 2, 2), Value::const_(0, 2))
                .simplify();
        assert_eq!(*state.get_value(&out).expect("out should be set"), expected);
    }

    #[test]
    fn join_indirect_write_missing_in_self_top_behavior_becomes_top() {
        // One branch writes RAX to *[addr]:8; the other branch writes nothing there.
        // With MergeBehavior::Top, the merged result should be Top.
        let arch = test_arch();
        let addr_reg = reg(0x1000, 8);
        let rax = reg(0x2000, 8);

        let load_key = Value::load(Value::entry(addr_reg), 8, 1);

        let mut self_state = ValuationState::new_with_behavior(arch.clone(), MergeBehavior::Top);
        let mut other_state = ValuationState::new_with_behavior(arch, MergeBehavior::Top);
        other_state
            .valuation
            .add(load_key.clone(), Value::entry(rax));

        self_state.join(&other_state);

        assert_eq!(
            self_state.valuation.indirect_writes.get(&load_key),
            Some(&Value::Top),
        );
    }

    #[test]
    fn join_indirect_write_missing_in_self_choice_behavior_becomes_choice() {
        // One branch writes a constant to *[addr]:8; the other branch writes nothing there.
        // With MergeBehavior::Choice, the merged result should be Choice(load_key, written_value).
        let arch = test_arch();
        let addr_reg = reg(0x1000, 8);

        let load_key = Value::load(Value::entry(addr_reg), 8, 1);
        let written = Value::const_(0x42, 8);

        let mut self_state = ValuationState::new_with_behavior(arch.clone(), MergeBehavior::Choice);
        let mut other_state = ValuationState::new_with_behavior(arch, MergeBehavior::Choice);
        other_state.valuation.add(load_key.clone(), written.clone());

        self_state.join(&other_state);

        let result = self_state
            .valuation
            .indirect_writes
            .get(&load_key)
            .expect("should have an entry after join");
        let choice = result.as_choice().expect("expected a Choice node");
        assert!(
            (choice.0.as_ref() == &load_key && choice.1.as_ref() == &written)
                || (choice.1.as_ref() == &load_key && choice.0.as_ref() == &written),
            "expected Choice(load_key, written) in any order, got {result:?}",
        );
    }
}
