facet-reflect 0.44.4

use super::*;
use crate::HasFields;
#[cfg(feature = "std")]
use core::cell::RefCell;
use hashbrown::{HashMap, HashSet};

#[cfg(feature = "std")]
thread_local! {
    static INVARIANT_SUBTREE_CACHE: RefCell<HashMap<facet_core::ConstTypeId, bool>> =
        RefCell::new(HashMap::new());
}

fn shape_subtree_has_invariants(
    shape: &'static Shape,
    cache: &mut HashMap<facet_core::ConstTypeId, Option<bool>>,
) -> bool {
    #[cfg(feature = "std")]
    if let Some(cached) = INVARIANT_SUBTREE_CACHE.with(|memo| memo.borrow().get(&shape.id).copied())
    {
        return cached;
    }

    if let Some(cached) = cache.get(&shape.id) {
        // `None` means we're currently evaluating this shape in a recursive cycle.
        // Returning false here breaks recursion; direct invariants are checked before descent.
        return cached.unwrap_or(false);
    }

    cache.insert(shape.id, None);

    let has_invariants = if shape.vtable.has_invariants() {
        true
    } else {
        match shape.ty {
            Type::User(UserType::Struct(struct_ty)) => struct_ty
                .fields
                .iter()
                .any(|field| shape_subtree_has_invariants(field.shape.get(), cache)),
            Type::User(UserType::Enum(enum_ty)) => enum_ty.variants.iter().any(|variant| {
                variant
                    .data
                    .fields
                    .iter()
                    .any(|field| shape_subtree_has_invariants(field.shape.get(), cache))
            }),
            _ => match shape.def {
                Def::List(list) => shape_subtree_has_invariants(list.t(), cache),
                Def::Array(array) => shape_subtree_has_invariants(array.t(), cache),
                Def::Slice(slice) => shape_subtree_has_invariants(slice.t(), cache),
                Def::Map(map) => {
                    shape_subtree_has_invariants(map.k(), cache)
                        || shape_subtree_has_invariants(map.v(), cache)
                }
                Def::Set(set) => shape_subtree_has_invariants(set.t(), cache),
                Def::Option(opt) => shape_subtree_has_invariants(opt.t(), cache),
                Def::Result(result) => {
                    shape_subtree_has_invariants(result.t(), cache)
                        || shape_subtree_has_invariants(result.e(), cache)
                }
                Def::Pointer(ptr) => ptr
                    .pointee()
                    .is_some_and(|pointee| shape_subtree_has_invariants(pointee, cache)),
                _ => false,
            },
        }
    };

    #[cfg(feature = "std")]
    INVARIANT_SUBTREE_CACHE.with(|memo| {
        memo.borrow_mut().insert(shape.id, has_invariants);
    });
    cache.insert(shape.id, Some(has_invariants));
    has_invariants
}

fn validate_invariants_recursive<'mem, 'facet>(
    value: Peek<'mem, 'facet>,
    visited: &mut HashSet<crate::ValueId>,
    shape_cache: &mut HashMap<facet_core::ConstTypeId, Option<bool>>,
) -> Result<(), (&'static Shape, String)> {
    if !shape_subtree_has_invariants(value.shape(), shape_cache) {
        return Ok(());
    }

    let id = value.id();
    if !visited.insert(id) {
        return Ok(());
    }

    if let Some(result) = unsafe { value.shape().call_invariants(value.data()) }
        && let Err(message) = result
    {
        return Err((value.shape(), message));
    }

    match value.shape().ty {
        Type::User(UserType::Struct(_)) => {
            if let Ok(peek_struct) = value.into_struct() {
                for (field, child) in peek_struct.fields() {
                    if shape_subtree_has_invariants(field.shape.get(), shape_cache) {
                        validate_invariants_recursive(child, visited, shape_cache)?;
                    }
                }
            }
        }
        Type::User(UserType::Enum(_)) => {
            if let Ok(peek_enum) = value.into_enum() {
                for (field, child) in peek_enum.fields() {
                    if shape_subtree_has_invariants(field.shape.get(), shape_cache) {
                        validate_invariants_recursive(child, visited, shape_cache)?;
                    }
                }
            }
        }
        _ => match value.shape().def {
            Def::List(_) | Def::Array(_) | Def::Slice(_) => {
                if let Ok(list_like) = value.into_list_like()
                    && shape_subtree_has_invariants(list_like.def.t(), shape_cache)
                {
                    for elem in list_like.iter() {
                        validate_invariants_recursive(elem, visited, shape_cache)?;
                    }
                }
            }
            Def::Map(_) => {
                if let Ok(map) = value.into_map() {
                    let def = map.def();
                    let key_has_invariants = shape_subtree_has_invariants(def.k(), shape_cache);
                    let value_has_invariants = shape_subtree_has_invariants(def.v(), shape_cache);
                    if key_has_invariants || value_has_invariants {
                        for (key, val) in map.iter() {
                            if key_has_invariants {
                                validate_invariants_recursive(key, visited, shape_cache)?;
                            }
                            if value_has_invariants {
                                validate_invariants_recursive(val, visited, shape_cache)?;
                            }
                        }
                    }
                }
            }
            Def::Set(_) => {
                if let Ok(set) = value.into_set()
                    && shape_subtree_has_invariants(set.def().t(), shape_cache)
                {
                    for elem in set.iter() {
                        validate_invariants_recursive(elem, visited, shape_cache)?;
                    }
                }
            }
            Def::Option(_) => {
                if let Ok(opt) = value.into_option()
                    && let Some(inner) = opt.value()
                    && shape_subtree_has_invariants(inner.shape(), shape_cache)
                {
                    validate_invariants_recursive(inner, visited, shape_cache)?;
                }
            }
            Def::Result(_) => {
                if let Ok(result) = value.into_result() {
                    if let Some(ok) = result.ok() {
                        validate_invariants_recursive(ok, visited, shape_cache)?;
                    }
                    if let Some(err) = result.err() {
                        validate_invariants_recursive(err, visited, shape_cache)?;
                    }
                }
            }
            Def::Pointer(_) => {
                if let Ok(ptr) = value.into_pointer()
                    && let Some(inner) = ptr.borrow_inner()
                    && shape_subtree_has_invariants(inner.shape(), shape_cache)
                {
                    validate_invariants_recursive(inner, visited, shape_cache)?;
                }
            }
            _ => {}
        },
    }

    Ok(())
}

////////////////////////////////////////////////////////////////////////////////////////////////////
// Build
////////////////////////////////////////////////////////////////////////////////////////////////////
impl<'facet, const BORROW: bool> Partial<'facet, BORROW> {
    /// Builds the value, consuming the Partial.
    pub fn build(mut self) -> Result<HeapValue<'facet, BORROW>, ReflectError> {
        use crate::typeplan::TypePlanNodeKind;

        if self.frames().len() != 1 {
            return Err(self.err(ReflectErrorKind::InvariantViolation {
                invariant: "Partial::build() expects a single frame — call end() until that's the case",
            }));
        }

        // Try the optimized path using precomputed FieldInitPlan (includes validators)
        // Extract frame info first (borrows only self.mode)
        let frame_info = self.mode.stack().last().map(|frame| {
            let variant_idx = match &frame.tracker {
                Tracker::Enum { variant_idx, .. } => Some(*variant_idx),
                _ => None,
            };
            (frame.type_plan, variant_idx)
        });

        // Look up plans from the type plan node - need to resolve NodeId to get the actual node
        let plans_info = frame_info.and_then(|(type_plan_id, variant_idx)| {
            let type_plan = self.root_plan.node(type_plan_id);
            match &type_plan.kind {
                TypePlanNodeKind::Struct(struct_plan) => Some(struct_plan.fields),
                TypePlanNodeKind::Enum(enum_plan) => {
                    let variants = self.root_plan.variants(enum_plan.variants);
                    variant_idx.and_then(|idx| variants.get(idx).map(|v| v.fields))
                }
                _ => None,
            }
        });

        if let Some(plans_range) = plans_info {
            // Resolve the SliceRange to an actual slice
            let plans = self.root_plan.fields(plans_range);
            // Now mutably borrow mode.stack to get the frame
            // (root_plan borrow of `plans` is still active but that's fine -
            // mode and root_plan are separate fields)
            let frame = self.mode.stack_mut().last_mut().unwrap();
            crate::trace!(
                "build(): Using optimized fill_and_require_fields for {}, tracker={:?}",
                frame.allocated.shape(),
                frame.tracker.kind()
            );
            frame
                .fill_and_require_fields(plans, plans.len(), &self.root_plan)
                .map_err(|e| self.err(e))?;
        } else {
            // Fall back to the old path if optimized path wasn't available
            let frame = self.frames_mut().last_mut().unwrap();
            crate::trace!(
                "build(): calling fill_defaults for {}, tracker={:?}, is_init={}",
                frame.allocated.shape(),
                frame.tracker.kind(),
                frame.is_init
            );
            if let Err(e) = frame.fill_defaults() {
                return Err(self.err(e));
            }
            crate::trace!(
                "build(): after fill_defaults, tracker={:?}, is_init={}",
                frame.tracker.kind(),
                frame.is_init
            );

            let frame = self.frames_mut().last_mut().unwrap();
            crate::trace!(
                "build(): calling require_full_initialization, tracker={:?}",
                frame.tracker.kind()
            );
            let result = frame.require_full_initialization();
            crate::trace!(
                "build(): require_full_initialization result: {:?}",
                result.is_ok()
            );
            result.map_err(|e| self.err(e))?
        }

        let frame = self.frames_mut().pop().unwrap();

        // Validate invariants on the full value tree (root + nested values).
        // Safety: the value is fully initialized at this point.
        let value_ptr = unsafe { frame.data.assume_init().as_const() };
        let root = unsafe { Peek::unchecked_new(value_ptr, frame.allocated.shape()) };
        let mut visited = HashSet::new();
        let mut shape_cache = HashMap::new();
        if let Err((shape, message)) =
            validate_invariants_recursive(root, &mut visited, &mut shape_cache)
        {
            // Put the frame back so Drop can handle cleanup properly
            self.frames_mut().push(frame);
            return Err(self.err(ReflectErrorKind::UserInvariantFailed { message, shape }));
        }

        // Mark as built to prevent Drop from cleaning up the value
        self.state = PartialState::Built;

        match frame
            .allocated
            .shape()
            .layout
            .sized_layout()
            .map_err(|_layout_err| {
                self.err(ReflectErrorKind::Unsized {
                    shape: frame.allocated.shape(),
                    operation: "build (final check for sized layout)",
                })
            }) {
            Ok(layout) => {
                // Determine if we should deallocate based on ownership
                let should_dealloc = frame.ownership.needs_dealloc();

                Ok(HeapValue {
                    guard: Some(Guard {
                        ptr: unsafe { NonNull::new_unchecked(frame.data.as_mut_byte_ptr()) },
                        layout,
                        should_dealloc,
                    }),
                    shape: frame.allocated.shape(),
                    phantom: PhantomData,
                })
            }
            Err(e) => {
                // Put the frame back for proper cleanup
                self.frames_mut().push(frame);
                Err(e)
            }
        }
    }

    /// Finishes deserialization in-place, validating the value without moving it.
    ///
    /// This is intended for use with [`from_raw`](Self::from_raw) where the value
    /// is deserialized into caller-provided memory (e.g., a `MaybeUninit<T>` on the stack).
    ///
    /// On success, the caller can safely assume the memory contains a fully initialized,
    /// valid value and call `MaybeUninit::assume_init()`.
    ///
    /// On failure, any partially initialized data is cleaned up (dropped), and the
    /// memory should be considered uninitialized.
    ///
    /// # Panics
    ///
    /// Panics if called with more than one frame on the stack (i.e., if you haven't
    /// called `end()` enough times to return to the root level).
    ///
    /// # Example
    ///
    /// ```ignore
    /// use std::mem::MaybeUninit;
    /// use facet_core::{Facet, PtrUninit};
    /// use facet_reflect::Partial;
    ///
    /// let mut slot = MaybeUninit::<MyStruct>::uninit();
    /// let ptr = PtrUninit::new(slot.as_mut_ptr().cast());
    ///
    /// let partial = unsafe { Partial::from_raw_with_shape(ptr, MyStruct::SHAPE)? };
    /// // ... deserialize into partial ...
    /// partial.finish_in_place()?;
    ///
    /// // Now safe to assume initialized
    /// let value = unsafe { slot.assume_init() };
    /// ```
    pub fn finish_in_place(mut self) -> Result<(), ReflectError> {
        use crate::typeplan::TypePlanNodeKind;

        if self.frames().len() != 1 {
            return Err(self.err(ReflectErrorKind::InvariantViolation {
                invariant: "Partial::finish_in_place() expects a single frame — call end() until that's the case",
            }));
        }

        // Try the optimized path using precomputed FieldInitPlan (includes validators)
        // Extract frame info first (borrows only self.mode)
        let frame_info = self.mode.stack().last().map(|frame| {
            let variant_idx = match &frame.tracker {
                Tracker::Enum { variant_idx, .. } => Some(*variant_idx),
                _ => None,
            };
            (frame.type_plan, variant_idx)
        });

        // Look up plans from the type plan node - need to resolve NodeId to get the actual node
        let plans_info = frame_info.and_then(|(type_plan_id, variant_idx)| {
            let type_plan = self.root_plan.node(type_plan_id);
            match &type_plan.kind {
                TypePlanNodeKind::Struct(struct_plan) => Some(struct_plan.fields),
                TypePlanNodeKind::Enum(enum_plan) => {
                    let variants = self.root_plan.variants(enum_plan.variants);
                    variant_idx.and_then(|idx| variants.get(idx).map(|v| v.fields))
                }
                _ => None,
            }
        });

        if let Some(plans_range) = plans_info {
            // Resolve the SliceRange to an actual slice
            let plans = self.root_plan.fields(plans_range);
            // Now mutably borrow mode.stack to get the frame
            // (root_plan borrow of `plans` is still active but that's fine -
            // mode and root_plan are separate fields)
            let frame = self.mode.stack_mut().last_mut().unwrap();
            crate::trace!(
                "finish_in_place(): Using optimized fill_and_require_fields for {}, tracker={:?}",
                frame.allocated.shape(),
                frame.tracker.kind()
            );
            frame
                .fill_and_require_fields(plans, plans.len(), &self.root_plan)
                .map_err(|e| self.err(e))?;
        } else {
            // Fall back to the old path if optimized path wasn't available
            let frame = self.frames_mut().last_mut().unwrap();
            crate::trace!(
                "finish_in_place(): calling fill_defaults for {}, tracker={:?}, is_init={}",
                frame.allocated.shape(),
                frame.tracker.kind(),
                frame.is_init
            );
            if let Err(e) = frame.fill_defaults() {
                return Err(self.err(e));
            }
            crate::trace!(
                "finish_in_place(): after fill_defaults, tracker={:?}, is_init={}",
                frame.tracker.kind(),
                frame.is_init
            );

            let frame = self.frames_mut().last_mut().unwrap();
            crate::trace!(
                "finish_in_place(): calling require_full_initialization, tracker={:?}",
                frame.tracker.kind()
            );
            let result = frame.require_full_initialization();
            crate::trace!(
                "finish_in_place(): require_full_initialization result: {:?}",
                result.is_ok()
            );
            result.map_err(|e| self.err(e))?
        }

        let frame = self.frames_mut().pop().unwrap();

        // Validate invariants on the full value tree (root + nested values).
        // Safety: the value is fully initialized at this point.
        let value_ptr = unsafe { frame.data.assume_init().as_const() };
        let root = unsafe { Peek::unchecked_new(value_ptr, frame.allocated.shape()) };
        let mut visited = HashSet::new();
        let mut shape_cache = HashMap::new();
        if let Err((shape, message)) =
            validate_invariants_recursive(root, &mut visited, &mut shape_cache)
        {
            // Put the frame back so Drop can handle cleanup properly
            self.frames_mut().push(frame);
            return Err(self.err(ReflectErrorKind::UserInvariantFailed { message, shape }));
        }

        // Mark as built to prevent Drop from cleaning up the now-valid value.
        // The caller owns the memory and will handle the value from here.
        self.state = PartialState::Built;

        // Frame is dropped here without deallocation (External ownership doesn't dealloc)
        Ok(())
    }
}